update function signature

Author: jgbradley1

backend/manage-indexing-jobs.py

@@ -109,7 +109,7 @@ def main():
     azure_storage_client_manager = AzureClientManager()
     job_container_store_client = (
         azure_storage_client_manager.get_cosmos_container_client(
-            database_name="graphrag", container_name="jobs"
+            database="graphrag", container="jobs"
         )
     )
     # retrieve status of all index jobs that are scheduled or running

backend/src/api/azure_clients.py

@@ -174,7 +174,7 @@ def get_cosmos_database_client(self, database_name: str) -> DatabaseProxy:
         return self._cosmos_client.get_database_client(database=database_name)
 
     def get_cosmos_container_client(
-        self, database_name: str, container_name: str
+        self, database: str, container: str
     ) -> ContainerProxy:
         """
         Returns a Cosmos container client.
@@ -187,5 +187,5 @@ def get_cosmos_container_client(
             ContainerProxy: The Cosmos DB container client.
         """
         return self._cosmos_client.get_database_client(
-            database=database_name
-        ).get_container_client(container=container_name)
+            database=database
+        ).get_container_client(container=container)

backend/src/api/common.py

@@ -44,12 +44,10 @@ def delete_cosmos_container_item(container: str, item_id: str):
     If exception is raised, the calling function should catch it.
     """
     azure_client_manager = AzureClientManager()
-    # cosmos_client = azure_client_manager.get_cosmos_client()
     try:
         azure_client_manager.get_cosmos_container_client(
-            "graphrag", container
+            database="graphrag", container=container
         ).delete_item(item_id, item_id)
-        # cosmos_client.get_database_client(database).get_container_client(container).delete_item(item_id, item_id)
     except exceptions.CosmosResourceNotFoundError:
         # If item does not exist, do nothing
         pass
@@ -74,7 +72,7 @@ def validate_index_file_exist(index_name: str, file_name: str):
     azure_client_manager = AzureClientManager()
     try:
         cosmos_container_client = azure_client_manager.get_cosmos_container_client(
-            database_name="graphrag", container_name="container-store"
+            database="graphrag", container="container-store"
         )
         cosmos_container_client.read_item(index_name, index_name)
     except Exception:
@@ -176,7 +174,7 @@ def retrieve_original_blob_container_name(sanitized_name: str) -> str | None:
     azure_client_manager = AzureClientManager()
     try:
         container_store_client = azure_client_manager.get_cosmos_container_client(
-            database_name="graphrag", container_name="container-store"
+            database="graphrag", container="container-store"
         )
         for item in container_store_client.read_all_items():
             if item["id"] == sanitized_name:

backend/src/api/data.py

@@ -46,7 +46,7 @@ async def get_all_data_storage_containers():
     items = []
     try:
         container_store_client = azure_client_manager.get_cosmos_container_client(
-            database_name="graphrag", container_name="container-store"
+            database="graphrag", container="container-store"
         )
         for item in container_store_client.read_all_items():
             if item["type"] == "data":
@@ -161,7 +161,7 @@ async def upload_files(
             await asyncio.gather(*tasks)
         # update container-store in cosmosDB since upload process was successful
         container_store_client = azure_client_manager.get_cosmos_container_client(
-            database_name="graphrag", container_name="container-store"
+            database="graphrag", container="container-store"
         )
         container_store_client.upsert_item({
             "id": sanitized_storage_name,

backend/src/api/index.py

@@ -162,7 +162,7 @@ async def _start_indexing_pipeline(index_name: str):
         blob_service_client.create_container(sanitized_index_name)
 
     cosmos_container_client = azure_client_manager.get_cosmos_container_client(
-        database_name="graphrag", container_name="container-store"
+        database="graphrag", container="container-store"
     )
     cosmos_container_client.upsert_item({
         "id": sanitized_index_name,
@@ -331,7 +331,7 @@ async def get_all_indexes():
     try:
         azure_client_manager = AzureClientManager()
         container_store_client = azure_client_manager.get_cosmos_container_client(
-            database_name="graphrag", container_name="container-store"
+            database="graphrag", container="container-store"
         )
         for item in container_store_client.read_all_items():
             if item["type"] == "index":
@@ -413,7 +413,7 @@ async def delete_index(index_name: str):
         # update container-store in cosmosDB
         try:
             container_store_client = azure_client_manager.get_cosmos_container_client(
-                database_name="graphrag", container_name="container-store"
+                database="graphrag", container="container-store"
             )
             container_store_client.delete_item(
                 item=sanitized_index_name, partition_key=sanitized_index_name
@@ -424,7 +424,7 @@ async def delete_index(index_name: str):
         # update jobs database in cosmosDB
         try:
             jobs_container = azure_client_manager.get_cosmos_container_client(
-                database_name="graphrag", container_name="jobs"
+                database="graphrag", container="jobs"
             )
             jobs_container.delete_item(
                 item=sanitized_index_name, partition_key=sanitized_index_name

backend/src/models.py

@@ -111,7 +111,7 @@ class PipelineJob:
     def _jobs_container():
         azure_storage_client = AzureClientManager()
         return azure_storage_client.get_cosmos_container_client(
-            database_name="graphrag", container_name="jobs"
+            database="graphrag", container="jobs"
         )
 
     @classmethod