Update the utilities and relevant wrappers to user the KeystoneApi Cl…

…ient Class (#243)

* update crc_proposal_end to use client class

* small fixes for getting login to work

* changes after testing on cluster

* tie other utility functions into using the client class

* update the other wrappers to use client

* typo in endpoint for getting clusters

apps/crc_proposal_end.py

@@ -33,16 +33,16 @@ def app_logic(self, args: Namespace) -> None:
         """
 
         Slurm.check_slurm_account_exists(args.account)
-        auth_header = get_auth_header(KEYSTONE_URL,
-                                      {'username': os.environ["USER"],
-                                       'password': getpass("Please enter your CRC login password:\n")})
-        keystone_group_id = get_researchgroup_id(KEYSTONE_URL, args.account, auth_header)
-        alloc_requests = get_active_requests(KEYSTONE_URL, keystone_group_id, auth_header)
+        keystone_session = KeystoneApi()
+        keystone_session.login(username=os.environ["USER"], password=getpass("Please enter your CRC login password:\n"))
+
+        group_id = get_researchgroup_id(keystone_session, args.account)
+        alloc_requests = get_active_requests(keystone_session, group_id)
 
         if not alloc_requests:
             print(f"\033[91m\033[1mNo active allocation information found in accounting system for '{args.account}'!\n")
             print("Showing end date for most recently expired Resource Allocation Request:\033[0m")
-            alloc_requests = get_most_recent_expired_request(KEYSTONE_URL, keystone_group_id, auth_header)
+            alloc_requests = get_most_recent_expired_request(keystone_session, group_id)
 
         for request in alloc_requests:
             print(f"'{request['title']}' ends on {request['expire']} ")

apps/crc_sus.py

@@ -58,21 +58,19 @@ def app_logic(self, args: Namespace) -> None:
         """
 
         Slurm.check_slurm_account_exists(account_name=args.account)
-        auth_header = get_auth_header(KEYSTONE_URL,
-                                      {'username': os.environ["USER"],
-                                       'password': getpass("Please enter your CRC login password:\n")})
-        # Determine if provided or default account is in Keystone
-        keystone_group_id = get_researchgroup_id(KEYSTONE_URL, args.account, auth_header)
-        alloc_requests = get_active_requests(KEYSTONE_URL, keystone_group_id, auth_header)
+        keystone_session = KeystoneApi()
+        keystone_session.login(username=os.environ["USER"], password=getpass("Please enter your CRC login password:\n"))
+
+        group_id = get_researchgroup_id(keystone_session, args.account)
+        alloc_requests = get_active_requests(keystone_session, group_id)
+
         if not alloc_requests:
             print(f"\033[91m\033[1mNo active allocation information found in accounting system for '{args.account}'!\n")
-            print("Showing SUs for most recently expired Resource Allocation Request:\033[0m")
-            alloc_requests = get_most_recent_expired_request(KEYSTONE_URL, keystone_group_id, auth_header)
-
-        per_cluster_totals = get_per_cluster_totals(alloc_requests,
-                                                    get_enabled_cluster_ids(KEYSTONE_URL, auth_header),
-                                                    auth_header)
+            print("Showing end date for most recently expired Resource Allocation Request:\033[0m")
+            alloc_requests = get_most_recent_expired_request(keystone_session, group_id)
 
+        per_cluster_totals = get_per_cluster_totals(keystone_session, alloc_requests,
+                                                    get_enabled_cluster_ids(keystone_session))
         earliest_date = get_earliest_startdate(alloc_requests)
 
         for cluster in per_cluster_totals:

apps/crc_usage.py

@@ -96,25 +96,24 @@ def app_logic(self, args: Namespace) -> None:
         """
 
         Slurm.check_slurm_account_exists(account_name=args.account)
-        auth_header = get_auth_header(KEYSTONE_URL,
-                                      {'username': os.environ["USER"],
-                                       'password': getpass("Please enter your CRC login password:\n")})
+        keystone_session = KeystoneApi()
+        keystone_session.login(username=os.environ["USER"], password=getpass("Please enter your CRC login password:\n"))
 
         # Gather AllocationRequests from Keystone
-        keystone_group_id = get_researchgroup_id(KEYSTONE_URL, args.account, auth_header)
+        group_id = get_researchgroup_id(keystone_session, args.account)
+        alloc_requests = get_active_requests(keystone_session, group_id)
 
-        alloc_requests = get_active_requests(KEYSTONE_URL, keystone_group_id, auth_header)
         if not alloc_requests:
             print(f"\033[91m\033[1mNo active allocation information found in accounting system for '{args.account}'!\n")
             print("Showing usage information for most recently expired Resource Allocation Request: \033[0m")
-            alloc_requests = get_most_recent_expired_request(KEYSTONE_URL, keystone_group_id, auth_header)
+            alloc_requests = get_most_recent_expired_request(keystone_session, group_id)
 
-        clusters = get_enabled_cluster_ids(KEYSTONE_URL, auth_header)
+        clusters = get_enabled_cluster_ids(keystone_session)
 
         self.print_summary_table(alloc_requests,
                                  args.account,
-                                 get_per_cluster_totals(alloc_requests, clusters, auth_header, per_request=True))
+                                 get_per_cluster_totals(keystone_session, alloc_requests, clusters, per_request=True))
 
         self.print_usage_table(args.account,
-                               get_per_cluster_totals(alloc_requests, clusters, auth_header),
+                               get_per_cluster_totals(keystone_session, alloc_requests, clusters),
                                get_earliest_startdate(alloc_requests))

apps/utils/keystone.py

@@ -11,7 +11,7 @@
 
 # Default API configuratipn
 KEYSTONE_URL = "https://keystone.crc.pitt.edu"
-KEYSTONE_AUTH_ENDPOINT = 'authentication/new'
+KEYSTONE_AUTH_ENDPOINT = 'authentication/new/'
 RAWUSAGE_RESET_DATE = date.fromisoformat('2024-05-07')
 
 
@@ -47,7 +47,7 @@ def login(self, username: str, password: str, endpoint: str = KEYSTONE_AUTH_ENDP
             timeout=self._timeout
         )
         response.raise_for_status()
-        self._token = response.json().get("token")
+        self._token = response.json().get("access")
 
     def _get_headers(self) -> Dict[str, str]:
         """Constructs the headers for an authenticated request.
@@ -216,56 +216,39 @@ def delete(self, endpoint: str, response_type: ResponseContentType = 'json') ->
         return self._process_response(response, response_type)
 
 
-def get_auth_header(keystone_url: str, auth_header: dict) -> dict:
-    """ Generate an authorization header to be used for accessing information from keystone"""
-
-    response = requests.post(f"{keystone_url}/authentication/new/", json=auth_header, timeout=10)
-    response.raise_for_status()
-    tokens = response.json()
-    return {"Authorization": f"Bearer {tokens['access']}"}
-
-
-def get_request_allocations(keystone_url: str, request_pk: int, auth_header: dict) -> dict:
+def get_request_allocations(keystone_client: KeystoneApi, request_pk: int) -> dict:
     """Get All Allocation information from keystone for a given request"""
 
-    response = requests.get(f"{keystone_url}/allocations/allocations/?request={request_pk}",
-                            headers=auth_header,
-                            timeout=10
-                            )
-    response.raise_for_status()
-    return response.json()
+    return keystone_client.get('allocations/allocations/', {'request': request_pk}, 'json')
 
 
-def get_active_requests(keystone_url: str, group_pk: int, auth_header: dict) -> [dict]:
+def get_active_requests(keystone_client: KeystoneApi, group_pk: int) -> [dict]:
     """Get all active AllocationRequest information from keystone for a given group"""
 
     today = date.today().isoformat()
-    response = requests.get(
-        f"{keystone_url}/allocations/requests/?group={group_pk}&status=AP&active__lte={today}&expire__gt={today}",
-        headers=auth_header,
-        timeout=10
-    )
-    response.raise_for_status()
-    return [request for request in response.json()]
+    return [request for request in keystone_client.get('allocations/requests/',
+                                                       {'group': group_pk,
+                                                        'status': 'AP',
+                                                        'active__lte': today,
+                                                        'expire__gt': today},
+                                                       'json'
+                                                       )]
 
 
-def get_researchgroup_id(keystone_url: str, account_name: str, auth_header: dict) -> int:
+def get_researchgroup_id(keystone_client: KeystoneApi, account_name: str) -> int:
     """Get the Researchgroup ID from keystone for the specified Slurm account"""
 
-    response = requests.get(f"{keystone_url}/users/researchgroups/?name={account_name}",
-                            headers=auth_header,
-                            timeout=10
-                            )
-    response.raise_for_status()
-
+    # Attempt to get the primary key for the ResearchGroup
     try:
-        group_id = int(response.json()[0]['id'])
+        keystone_group_id = keystone_client.get('users/researchgroups/',
+                                                {'name': account_name},
+                                                'json')[0]['id']
     except IndexError:
         print(f"No Slurm Account found in the accounting system for '{account_name}'. \n"
               f"Please submit a ticket to the CRC team to ensure your allocation was properly configured")
         exit()
 
-    return group_id
+    return keystone_group_id
 
 
 def get_earliest_startdate(alloc_requests: [dict]) -> date:
@@ -281,43 +264,40 @@ def get_earliest_startdate(alloc_requests: [dict]) -> date:
     return max(earliest_date, RAWUSAGE_RESET_DATE)
 
 
-def get_most_recent_expired_request(keystone_url: str, group_pk: int, auth_header: dict) -> [dict]:
+def get_most_recent_expired_request(keystone_client: KeystoneApi, group_pk: int) -> [dict]:
     """Get the single most recently expired AllocationRequest information from keystone for a given group"""
 
     today = date.today().isoformat()
-    response = requests.get(
-        f"{keystone_url}/allocations/requests/?ordering=-expire&group={group_pk}&status=AP&expire__lte={today}",
-        headers=auth_header,
-        timeout=10
-    )
-    response.raise_for_status()
-
-    return [response.json()[0]]
+    return [keystone_client.get('allocations/requests/',
+                                {'group': group_pk,
+                                 'status': 'AP',
+                                 'ordering': '-expire',
+                                 'expire__lte': today},
+                                'json'
+                                )[0]]
 
 
-def get_enabled_cluster_ids(keystone_url: str, auth_header: dict) -> dict():
+def get_enabled_cluster_ids(keystone_client: KeystoneApi) -> dict():
     """Get the list of enabled clusters defined in Keystone along with their IDs"""
 
-    response = requests.get(f"{keystone_url}/allocations/clusters/?enabled=True", headers=auth_header, timeout=10)
-    response.raise_for_status()
     clusters = {}
-    for cluster in response.json():
+    for cluster in keystone_client.get('allocations/clusters/', {'enabled': True}, 'json'):
         clusters[cluster['id']] = cluster['name']
 
     return clusters
 
 
-def get_per_cluster_totals(alloc_requests: [dict],
+def get_per_cluster_totals(keystone_client: KeystoneApi,
+                           alloc_requests: [dict],
                            clusters: dict,
-                           auth_header: dict,
                            per_request: bool = False) -> dict:
     """Gather the awarded totals across the given requests on each cluster into a dictionary"""
 
     per_cluster_totals = {}
     for request in alloc_requests:
         if per_request:
             per_cluster_totals[request['id']] = {}
-        for allocation in get_request_allocations(KEYSTONE_URL, request['id'], auth_header):
+        for allocation in get_request_allocations(keystone_client, request['id']):
             cluster = clusters[allocation['cluster']]
             if per_request:
                 per_cluster_totals[request['id']].setdefault(cluster, 0)