From 61e2a390075c86f3f9cc2fc030de15f8393e8ca4 Mon Sep 17 00:00:00 2001
From: Varun Valada <varun.valada@gmail.com>
Date: Thu, 19 Dec 2024 18:07:27 -0600
Subject: [PATCH] Update credential script to add ability to set maximum
 reservation time

---
 server/src/tools/client_credentials_admin.py | 27 ++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/server/src/tools/client_credentials_admin.py b/server/src/tools/client_credentials_admin.py
index 63d4feac..03366c36 100644
--- a/server/src/tools/client_credentials_admin.py
+++ b/server/src/tools/client_credentials_admin.py
@@ -87,6 +87,28 @@ def handle_allowed_queues_input() -> list:
     return allowed_queues
 
 
+def handle_max_reservation_input() -> list:
+    """Gets a user inputted max_reservation_time dict"""
+    max_reservation = {}
+    while True:
+        max_reservation_input = input(
+            "Do you want to specify a queue to set max reservation time for? (y/n): "
+        )
+        if max_reservation_input == "y":
+            queue_name = input(
+                "Enter the name of the queue(* to apply max reservation for"
+                + " all queues): "
+            )
+            max_reservation[queue_name] = int(
+                input("Enter the maximum reservation time for this queue: ")
+            )
+        elif max_reservation_input == "n":
+            break
+        else:
+            print("Invalid input. Please enter 'y' or 'n'.")
+    return max_reservation
+
+
 def confirm_dialogue(entry: dict) -> bool:
     """Prompts the user to confirm their changes"""
     print(entry)
@@ -136,11 +158,13 @@ def create_client_credential(db):
 
     max_priority = handle_max_priority_input()
     allowed_queues = handle_allowed_queues_input()
+    max_reservation = handle_max_reservation_input()
 
     entry = {
         "client_id": client_id,
         "max_priority": max_priority,
         "allowed_queues": allowed_queues,
+        "max_reservation_time": max_reservation,
     }
     if confirm_dialogue(entry):
         entry["client_secret_hash"] = client_secret_hash
@@ -161,11 +185,13 @@ def edit_client_credential(db):
 
     max_priority = handle_max_priority_input()
     allowed_queues = handle_allowed_queues_input()
+    max_reservation = handle_max_reservation_input()
     confirm_output = confirm_dialogue(
         {
             "client_id": client_id,
             "max_priority": max_priority,
             "allowed_queues": allowed_queues,
+            "max_reservation_time": max_reservation,
         }
     )
     if confirm_output:
@@ -175,6 +201,7 @@ def edit_client_credential(db):
                 "$set": {
                     "max_priority": max_priority,
                     "allowed_queues": allowed_queues,
+                    "max_reservation_time": max_reservation,
                 }
             },
             return_document=ReturnDocument.AFTER,