initial cleanup of code to support integration and unit test via pytest

Author: jgbradley1

.devcontainer/Dockerfile

@@ -1,4 +1,4 @@
-# Fmore information about the base image visit:
+# For more information about the base image visit:
 # https://mcr.microsoft.com/en-us/artifact/mar/devcontainers/python/about
 FROM mcr.microsoft.com/devcontainers/python:3.10-bookworm
 

backend/.coveragerc

@@ -0,0 +1,6 @@
+[run]
+omit =
+    **/__init__.py
+    src/api/query*
+    src/utils/query.py
+    src/models.py

backend/manage-indexing-jobs.py

@@ -19,8 +19,8 @@
 
 from src.api.azure_clients import AzureClientManager
 from src.api.common import sanitize_name
+from src.logger.logger_singleton import LoggerSingleton
 from src.models import PipelineJob
-from src.reporting.reporter_singleton import ReporterSingleton
 from src.typing.pipeline import PipelineJobState
 
 
@@ -47,7 +47,7 @@ def schedule_indexing_job(index_name: str):
             body=job_manifest, namespace=os.environ["AKS_NAMESPACE"]
         )
     except Exception:
-        reporter = ReporterSingleton().get_instance()
+        reporter = LoggerSingleton().get_instance()
         reporter.on_error(
             "Index job manager encountered error scheduling indexing job",
         )

backend/poetry.lock

@@ -4,6 +4,6 @@ asyncio_default_fixture_loop_scope="function"
 asyncio_mode=auto
 ; use well known credentials for Cosmos DB emulator and Azure Storage emulator
 env =
-    COSMOS_CONNECTION_STRING=AccountEndpoint=http://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==
+    COSMOS_CONNECTION_STRING=AccountEndpoint=http://127.0.0.1:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==
     STORAGE_CONNECTION_STRING=DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://127.0.0.1:10000/devstoreaccount1;
     TESTING=1

backend/pytest.ini

@@ -11,12 +11,11 @@
 from azure.identity import DefaultAzureCredential
 from azure.storage.blob import BlobServiceClient
 from azure.storage.blob.aio import BlobServiceClient as BlobServiceClientAsync
-from environs import Env
 
 ENDPOINT_ERROR_MSG = "Could not find connection string in environment variables"
 
 
-class CosmosClientSingleton:
+class _CosmosClientSingleton:
     """
     Singleton class for a CosmosClient instance.
 
@@ -29,19 +28,19 @@ class CosmosClientSingleton:
     @classmethod
     def get_instance(cls):
         if not cls._instance:
-            endpoint = os.getenv("COSMOS_URI_ENDPOINT")
             conn_string = os.getenv("COSMOS_CONNECTION_STRING")
             if conn_string:
                 cls._instance = CosmosClient.from_connection_string(conn_string)
             else:
+                endpoint = os.getenv("COSMOS_URI_ENDPOINT")
                 credential = DefaultAzureCredential()
                 cls._instance = CosmosClient(endpoint, credential)
         return cls._instance
 
 
-class BlobServiceClientSingleton:
+class _BlobServiceClientSingleton:
     """
-    Singleton class for BlobServiceClient.
+    Singleton class for a BlobServiceClient instance.
 
     If a connection string is available, it will be used to create the BlobServiceClient instance.
     Otherwise assume managed identity is used.
@@ -61,28 +60,16 @@ def get_instance(cls) -> BlobServiceClient:
                 cls._instance = BlobServiceClient(account_url, credential=credential)
         return cls._instance
 
-    @classmethod
-    def get_storage_account_name(cls) -> str:
-        conn_string = os.getenv("STORAGE_CONNECTION_STRING")
-        if conn_string:
-            # parse account name from the connection string
-            meta_info = {}
-            for meta_data in conn_string.split(";"):
-                if not meta_data:
-                    continue
-                m = meta_data.split("=", 1)
-                if len(m) != 2:
-                    continue
-                meta_info[m[0]] = m[1]
-            return meta_info["AccountName"]
-        else:
-            account_url = os.getenv("STORAGE_ACCOUNT_BLOB_URL")
-            return account_url.split("//")[1].split(".")[0]
 
+class _BlobServiceClientSingletonAsync:
+    """
+    Singleton class for a BlobServiceClientAsync instance.
+
+    If a connection string is available, it will be used to create the BlobServiceClientAsync instance.
+    Otherwise assume managed identity is used.
+    """
 
-class BlobServiceClientSingletonAsync:
     _instance = None
-    _env = Env()
 
     @classmethod
     def get_instance(cls) -> BlobServiceClientAsync:
@@ -100,43 +87,15 @@ def get_instance(cls) -> BlobServiceClientAsync:
                 )
         return cls._instance
 
-    @classmethod
-    def get_storage_account_name(cls) -> str:
-        conn_string = os.getenv("STORAGE_CONNECTION_STRING")
-        if conn_string:
-            # parse account name from the connection string
-            meta_info = {}
-            for meta_data in conn_string.split(";"):
-                if not meta_data:
-                    continue
-                m = meta_data.split("=", 1)
-                if len(m) != 2:
-                    continue
-                meta_info[m[0]] = m[1]
-            return meta_info["AccountName"]
-        else:
-            account_url = os.environ["STORAGE_ACCOUNT_BLOB_URL"]
-            return account_url.split("//")[1].split(".")[0]
-
-
-def get_database_client(database_name: str) -> DatabaseProxy:
-    client = CosmosClientSingleton.get_instance()
-    return client.get_database_client(database_name)
-
-
-def get_database_container_client(
-    database_name: str, container_name: str
-) -> ContainerProxy:
-    db_client = get_database_client(database_name)
-    return db_client.get_container_client(container_name)
-
 
 class AzureClientManager:
     """
     Manages the clients for Azure blob storage and Cosmos DB.
 
     Attributes:
-        azure_storage_blob_url (str): The blob endpoint for azure storage.
+        storage_blob_url (str): The blob endpoint for azure storage.
+        storage_account_name (str): The name of the azure storage account.
+        storage_account_hostname (str): The hostname of the azure blob storage account.
         cosmos_uri_endpoint (str): The uri endpoint for the Cosmos DB.
         _blob_service_client (BlobServiceClient): The blob service client.
         _blob_service_client_async (BlobServiceClientAsync): The asynchronous blob service client.
@@ -146,37 +105,34 @@ class AzureClientManager:
     """
 
     def __init__(self) -> None:
-        self.azure_storage_blob_url = os.getenv("STORAGE_ACCOUNT_BLOB_URL")
-        self.azure_storage_connection_string = os.getenv("STORAGE_CONNECTION_STRING")
+        self.storage_blob_url = os.getenv("STORAGE_ACCOUNT_BLOB_URL")
+        self.storage_connection_string = os.getenv("STORAGE_CONNECTION_STRING")
         self.cosmos_uri_endpoint = os.getenv("COSMOS_URI_ENDPOINT")
         self.cosmos_connection_string = os.getenv("COSMOS_CONNECTION_STRING")
-
-        if self.cosmos_connection_string:
-            self._cosmos_client = CosmosClient.from_connection_string(
-                self.cosmos_connection_string
-            )
-        else:
-            self._cosmos_client = CosmosClient(
-                self.cosmos_uri_endpoint, credential=DefaultAzureCredential()
-            )
-        if self.azure_storage_connection_string:
-            self._blob_service_client = BlobServiceClient.from_connection_string(
-                self.azure_storage_connection_string
-            )
-            self._blob_service_client_async = (
-                BlobServiceClientAsync.from_connection_string(
-                    self.azure_storage_connection_string
-                )
-            )
+        self._cosmos_client = _CosmosClientSingleton.get_instance()
+        self._blob_service_client = _BlobServiceClientSingleton.get_instance()
+        self._blob_service_client_async = (
+            _BlobServiceClientSingletonAsync.get_instance()
+        )
+
+        # parse account name from the azure storage connection string or blob url
+        if self.storage_connection_string:
+            meta_info = {}
+            for meta_data in self.storage_connection_string.split(";"):
+                if not meta_data:
+                    continue
+                m = meta_data.split("=", 1)
+                if len(m) != 2:
+                    continue
+                meta_info[m[0]] = m[1]
+            self.storage_account_name = meta_info["AccountName"]
         else:
-            self._blob_service_client = BlobServiceClient(
-                account_url=self.azure_storage_blob_url,
-                credential=DefaultAzureCredential(),
-            )
-            self._blob_service_client_async = BlobServiceClientAsync(
-                account_url=self.azure_storage_blob_url,
-                credential=DefaultAzureCredential(),
-            )
+            self.storage_account_name = self.storage_blob_url.split("//")[1].split(".")[
+                0
+            ]
+
+        # parse account hostname from the azure storage connection string or blob url
+        self.storage_account_hostname = self._blob_service_client.url.split("//")[1]
 
     def get_blob_service_client(self) -> BlobServiceClient:
         """
@@ -215,11 +171,7 @@ def get_cosmos_database_client(self, database_name: str) -> DatabaseProxy:
         Returns:
             DatabaseProxy: The Cosmos database client.
         """
-        if not hasattr(self, "_cosmos_database_client"):
-            self._cosmos_database_client = self._cosmos_client.get_database_client(
-                database=database_name
-            )
-        return self._cosmos_database_client
+        return self._cosmos_client.get_database_client(database=database_name)
 
     def get_cosmos_container_client(
         self, database_name: str, container_name: str
@@ -234,8 +186,6 @@ def get_cosmos_container_client(
         Returns:
             ContainerProxy: The Cosmos DB container client.
         """
-        if not hasattr(self, "_cosmos_container_client"):
-            self._cosmos_container_client = self.get_cosmos_database_client(
-                database_name=database_name
-            ).get_container_client(container=container_name)
-        return self._cosmos_container_client
+        return self._cosmos_client.get_database_client(
+            database=database_name
+        ).get_container_client(container=container_name)

backend/src/api/azure_clients.py

@@ -2,24 +2,37 @@
 # Licensed under the MIT License.
 
 import hashlib
+import os
 import re
 
+from azure.identity import DefaultAzureCredential
 from fastapi import HTTPException
 
-from src.api.azure_clients import (
-    AzureClientManager,
-    BlobServiceClientSingleton,
-)
+from src.api.azure_clients import AzureClientManager
 
-blob_service_client = BlobServiceClientSingleton.get_instance()
-azure_storage_client_manager = AzureClientManager()
+
+def get_pandas_storage_options() -> dict:
+    """Generate the storage options required by pandas to read parquet files from Storage."""
+    # For more information on the options available, see: https://github.com/fsspec/adlfs?tab=readme-ov-file#setting-credentials
+    azure_client_manager = AzureClientManager()
+    options = {
+        "account_name": azure_client_manager.storage_account_name,
+        "account_host": azure_client_manager.storage_account_hostname,
+    }
+    if os.getenv("STORAGE_CONNECTION_STRING"):
+        options["connection_string"] = os.getenv("STORAGE_CONNECTION_STRING")
+    else:
+        options["credential"] = DefaultAzureCredential()
+    return options
 
 
 def delete_blob_container(container_name: str):
     """
     Delete a blob container. If it does not exist, do nothing.
     If exception is raised, the calling function should catch it.
     """
+    azure_client_manager = AzureClientManager()
+    blob_service_client = azure_client_manager.get_blob_service_client()
     if blob_service_client.get_container_client(container_name).exists():
         blob_service_client.delete_container(container_name)
 
@@ -40,18 +53,18 @@ def validate_index_file_exist(index_name: str, file_name: str):
 
     Raises: ValueError
     """
-    # verify index_name is a valid index by checking container-store in cosmos db
+    azure_client_manager = AzureClientManager()
     try:
-        container_store_client = (
-            azure_storage_client_manager.get_cosmos_container_client(
-                database_name="graphrag", container_name="container-store"
-            )
+        cosmos_container_client = azure_client_manager.get_cosmos_container_client(
+            database_name="graphrag", container_name="container-store"
         )
-        container_store_client.read_item(index_name, index_name)
+        cosmos_container_client.read_item(index_name, index_name)
     except Exception:
         raise ValueError(f"Container {index_name} is not a valid index.")
     # check for file existence
-    index_container_client = blob_service_client.get_container_client(index_name)
+    index_container_client = (
+        azure_client_manager.get_blob_service_client().get_container_client(index_name)
+    )
     if not index_container_client.exists():
         raise ValueError(f"Container {index_name} not found.")
     if not index_container_client.get_blob_client(file_name).exists():
@@ -142,11 +155,10 @@ def retrieve_original_blob_container_name(sanitized_name: str) -> str | None:
     Returns: str
         The original human-readable name.
     """
+    azure_client_manager = AzureClientManager()
     try:
-        container_store_client = (
-            azure_storage_client_manager.get_cosmos_container_client(
-                database_name="graphrag", container_name="container-store"
-            )
+        container_store_client = azure_client_manager.get_cosmos_container_client(
+            database_name="graphrag", container_name="container-store"
         )
         for item in container_store_client.read_all_items():
             if item["id"] == sanitized_name:
@@ -156,31 +168,3 @@ def retrieve_original_blob_container_name(sanitized_name: str) -> str | None:
             status_code=500, detail="Error retrieving original blob name."
         )
     return None
-
-
-def retrieve_original_entity_config_name(sanitized_name: str) -> str | None:
-    """
-    Retrieve the original human-readable entity config name of a sanitized entity config name.
-
-    Args:
-    -----
-    sanitized_name (str)
-        The sanitized name to be converted back to the original name.
-
-    Returns: str
-        The original human-readable name.
-    """
-    try:
-        container_store_client = (
-            azure_storage_client_manager.get_cosmos_container_client(
-                database_name="graphrag", container_name="entities"
-            )
-        )
-        for item in container_store_client.read_all_items():
-            if item["id"] == sanitized_name:
-                return item["human_readable_name"]
-    except Exception:
-        raise HTTPException(
-            status_code=500, detail="Error retrieving original entity config name."
-        )
-    return None