Skip to content

Commit 778929a

Browse files
committed
update function signature
1 parent 1844fe4 commit 778929a

File tree

6 files changed

+14
-16
lines changed

6 files changed

+14
-16
lines changed

backend/manage-indexing-jobs.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ def main():
109109
azure_storage_client_manager = AzureClientManager()
110110
job_container_store_client = (
111111
azure_storage_client_manager.get_cosmos_container_client(
112-
database_name="graphrag", container_name="jobs"
112+
database="graphrag", container="jobs"
113113
)
114114
)
115115
# retrieve status of all index jobs that are scheduled or running

backend/src/api/azure_clients.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def get_cosmos_database_client(self, database_name: str) -> DatabaseProxy:
174174
return self._cosmos_client.get_database_client(database=database_name)
175175

176176
def get_cosmos_container_client(
177-
self, database_name: str, container_name: str
177+
self, database: str, container: str
178178
) -> ContainerProxy:
179179
"""
180180
Returns a Cosmos container client.
@@ -187,5 +187,5 @@ def get_cosmos_container_client(
187187
ContainerProxy: The Cosmos DB container client.
188188
"""
189189
return self._cosmos_client.get_database_client(
190-
database=database_name
191-
).get_container_client(container=container_name)
190+
database=database
191+
).get_container_client(container=container)

backend/src/api/common.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,10 @@ def delete_cosmos_container_item(container: str, item_id: str):
4444
If exception is raised, the calling function should catch it.
4545
"""
4646
azure_client_manager = AzureClientManager()
47-
# cosmos_client = azure_client_manager.get_cosmos_client()
4847
try:
4948
azure_client_manager.get_cosmos_container_client(
50-
"graphrag", container
49+
database="graphrag", container=container
5150
).delete_item(item_id, item_id)
52-
# cosmos_client.get_database_client(database).get_container_client(container).delete_item(item_id, item_id)
5351
except exceptions.CosmosResourceNotFoundError:
5452
# If item does not exist, do nothing
5553
pass
@@ -74,7 +72,7 @@ def validate_index_file_exist(index_name: str, file_name: str):
7472
azure_client_manager = AzureClientManager()
7573
try:
7674
cosmos_container_client = azure_client_manager.get_cosmos_container_client(
77-
database_name="graphrag", container_name="container-store"
75+
database="graphrag", container="container-store"
7876
)
7977
cosmos_container_client.read_item(index_name, index_name)
8078
except Exception:
@@ -176,7 +174,7 @@ def retrieve_original_blob_container_name(sanitized_name: str) -> str | None:
176174
azure_client_manager = AzureClientManager()
177175
try:
178176
container_store_client = azure_client_manager.get_cosmos_container_client(
179-
database_name="graphrag", container_name="container-store"
177+
database="graphrag", container="container-store"
180178
)
181179
for item in container_store_client.read_all_items():
182180
if item["id"] == sanitized_name:

backend/src/api/data.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ async def get_all_data_storage_containers():
4646
items = []
4747
try:
4848
container_store_client = azure_client_manager.get_cosmos_container_client(
49-
database_name="graphrag", container_name="container-store"
49+
database="graphrag", container="container-store"
5050
)
5151
for item in container_store_client.read_all_items():
5252
if item["type"] == "data":
@@ -161,7 +161,7 @@ async def upload_files(
161161
await asyncio.gather(*tasks)
162162
# update container-store in cosmosDB since upload process was successful
163163
container_store_client = azure_client_manager.get_cosmos_container_client(
164-
database_name="graphrag", container_name="container-store"
164+
database="graphrag", container="container-store"
165165
)
166166
container_store_client.upsert_item({
167167
"id": sanitized_storage_name,

backend/src/api/index.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@ async def _start_indexing_pipeline(index_name: str):
162162
blob_service_client.create_container(sanitized_index_name)
163163

164164
cosmos_container_client = azure_client_manager.get_cosmos_container_client(
165-
database_name="graphrag", container_name="container-store"
165+
database="graphrag", container="container-store"
166166
)
167167
cosmos_container_client.upsert_item({
168168
"id": sanitized_index_name,
@@ -331,7 +331,7 @@ async def get_all_indexes():
331331
try:
332332
azure_client_manager = AzureClientManager()
333333
container_store_client = azure_client_manager.get_cosmos_container_client(
334-
database_name="graphrag", container_name="container-store"
334+
database="graphrag", container="container-store"
335335
)
336336
for item in container_store_client.read_all_items():
337337
if item["type"] == "index":
@@ -413,7 +413,7 @@ async def delete_index(index_name: str):
413413
# update container-store in cosmosDB
414414
try:
415415
container_store_client = azure_client_manager.get_cosmos_container_client(
416-
database_name="graphrag", container_name="container-store"
416+
database="graphrag", container="container-store"
417417
)
418418
container_store_client.delete_item(
419419
item=sanitized_index_name, partition_key=sanitized_index_name
@@ -424,7 +424,7 @@ async def delete_index(index_name: str):
424424
# update jobs database in cosmosDB
425425
try:
426426
jobs_container = azure_client_manager.get_cosmos_container_client(
427-
database_name="graphrag", container_name="jobs"
427+
database="graphrag", container="jobs"
428428
)
429429
jobs_container.delete_item(
430430
item=sanitized_index_name, partition_key=sanitized_index_name

backend/src/models.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ class PipelineJob:
111111
def _jobs_container():
112112
azure_storage_client = AzureClientManager()
113113
return azure_storage_client.get_cosmos_container_client(
114-
database_name="graphrag", container_name="jobs"
114+
database="graphrag", container="jobs"
115115
)
116116

117117
@classmethod

0 commit comments

Comments
 (0)