Move to using managed identity for auth to CosmosDB. (#3806)

CHANGELOG.md

@@ -7,6 +7,7 @@ FEATURES:
 
 ENHANCEMENTS:
 * Switch from OpenCensus to OpenTelemetry for logging ([#3762](https://github.com/microsoft/AzureTRE/pull/3762))
+* Use managed identity for API connection to CosmosDB ([#345](https://github.com/microsoft/AzureTRE/issues/345))
 * Switch to Structured Firewall Logs ([#3816](https://github.com/microsoft/AzureTRE/pull/3816))
 
 BUG FIXES:

api_app/_version.py

@@ -1 +1 @@
-__version__ = "0.17.1"
+__version__ = "0.18.0"

api_app/api/dependencies/airlock.py

@@ -1,7 +1,7 @@
 from fastapi import Depends, HTTPException, Path, status
 from pydantic import UUID4
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.repositories.airlock_requests import AirlockRequestRepository
 from models.domain.airlock_request import AirlockRequest
 from db.errors import EntityDoesNotExist, UnableToAccessDatabase

api_app/api/dependencies/database.py

@@ -1,80 +1,86 @@
-from typing import Callable, Type
-
-from azure.cosmos.aio import CosmosClient
+from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy
 from azure.mgmt.cosmosdb.aio import CosmosDBManagementClient
-from fastapi import Depends, FastAPI, HTTPException
-from fastapi import Request, status
-from core import config, credentials
-from db.errors import UnableToAccessDatabase
-from db.repositories.base import BaseRepository
-from resources import strings
+
+from core.config import MANAGED_IDENTITY_CLIENT_ID, STATE_STORE_ENDPOINT, STATE_STORE_KEY, STATE_STORE_SSL_VERIFY, SUBSCRIPTION_ID, RESOURCE_MANAGER_ENDPOINT, CREDENTIAL_SCOPES, RESOURCE_GROUP_NAME, COSMOSDB_ACCOUNT_NAME, STATE_STORE_DATABASE
+from core.credentials import get_credential_async
 from services.logging import logger
 
 
-async def connect_to_db() -> CosmosClient:
-    logger.debug(f"Connecting to {config.STATE_STORE_ENDPOINT}")
+class Singleton(type):
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]
+
 
-    try:
-        async with credentials.get_credential_async() as credential:
-            primary_master_key = await get_store_key(credential)
+class Database(metaclass=Singleton):
 
-        if config.STATE_STORE_SSL_VERIFY:
+    _cosmos_client: CosmosClient = None
+    _database_proxy: DatabaseProxy = None
+
+    def __init__(cls):
+        pass
+
+    @classmethod
+    async def _connect_to_db(cls) -> CosmosClient:
+        logger.debug(f"Connecting to {STATE_STORE_ENDPOINT}")
+
+        credential = await get_credential_async()
+        if MANAGED_IDENTITY_CLIENT_ID:
+            logger.debug("Connecting with managed identity")
             cosmos_client = CosmosClient(
-                url=config.STATE_STORE_ENDPOINT, credential=primary_master_key
+                url=STATE_STORE_ENDPOINT,
+                credential=credential
             )
         else:
-            # ignore TLS (setup is a pain) when using local Cosmos emulator.
-            cosmos_client = CosmosClient(
-                config.STATE_STORE_ENDPOINT, primary_master_key, connection_verify=False
-            )
+            logger.debug("Connecting with key")
+            primary_master_key = await cls._get_store_key(credential)
+
+            if STATE_STORE_SSL_VERIFY:
+                logger.debug("Connecting with SSL verification")
+                cosmos_client = CosmosClient(
+                    url=STATE_STORE_ENDPOINT,
+                    credential=primary_master_key
+                )
+            else:
+                logger.debug("Connecting without SSL verification")
+                # ignore TLS (setup is a pain) when using local Cosmos emulator.
+                cosmos_client = CosmosClient(
+                    url=STATE_STORE_ENDPOINT,
+                    credential=primary_master_key,
+                    connection_verify=False
+                )
         logger.debug("Connection established")
         return cosmos_client
-    except Exception:
-        logger.exception("Connection to state store could not be established.")
-
-
-async def get_store_key(credential) -> str:
-    if config.STATE_STORE_KEY:
-        primary_master_key = config.STATE_STORE_KEY
-    else:
-        async with CosmosDBManagementClient(
-            credential,
-            subscription_id=config.SUBSCRIPTION_ID,
-            base_url=config.RESOURCE_MANAGER_ENDPOINT,
-            credential_scopes=config.CREDENTIAL_SCOPES
-        ) as cosmosdb_mng_client:
-            database_keys = await cosmosdb_mng_client.database_accounts.list_keys(
-                resource_group_name=config.RESOURCE_GROUP_NAME,
-                account_name=config.COSMOSDB_ACCOUNT_NAME,
-            )
-            primary_master_key = database_keys.primary_master_key
-
-    return primary_master_key
-
-
-async def get_db_client(app: FastAPI) -> CosmosClient:
-    if not hasattr(app.state, 'cosmos_client') or not app.state.cosmos_client:
-        app.state.cosmos_client = await connect_to_db()
-    return app.state.cosmos_client
 
+    @classmethod
+    async def _get_store_key(cls, credential) -> str:
+        logger.debug("Getting store key")
+        if STATE_STORE_KEY:
+            primary_master_key = STATE_STORE_KEY
+        else:
+            async with CosmosDBManagementClient(
+                credential,
+                subscription_id=SUBSCRIPTION_ID,
+                base_url=RESOURCE_MANAGER_ENDPOINT,
+                credential_scopes=CREDENTIAL_SCOPES
+            ) as cosmosdb_mng_client:
+                database_keys = await cosmosdb_mng_client.database_accounts.list_keys(
+                    resource_group_name=RESOURCE_GROUP_NAME,
+                    account_name=COSMOSDB_ACCOUNT_NAME,
+                )
+                primary_master_key = database_keys.primary_master_key
 
-async def get_db_client_from_request(request: Request) -> CosmosClient:
-    return await get_db_client(request.app)
+        return primary_master_key
 
+    @classmethod
+    async def get_container_proxy(cls, container_name) -> ContainerProxy:
+        if cls._cosmos_client is None:
+            cls._cosmos_client = await cls._connect_to_db()
 
-def get_repository(
-    repo_type: Type[BaseRepository],
-) -> Callable[[CosmosClient], BaseRepository]:
-    async def _get_repo(
-        client: CosmosClient = Depends(get_db_client_from_request),
-    ) -> BaseRepository:
-        try:
-            return await repo_type.create(client)
-        except UnableToAccessDatabase:
-            logger.exception(strings.STATE_STORE_ENDPOINT_NOT_RESPONDING)
-            raise HTTPException(
-                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
-                detail=strings.STATE_STORE_ENDPOINT_NOT_RESPONDING,
-            )
+        if cls._database_proxy is None:
+            cls._database_proxy = cls._cosmos_client.get_database_client(STATE_STORE_DATABASE)
 
-    return _get_repo
+        return cls._database_proxy.get_container_client(container_name)

api_app/api/dependencies/shared_services.py

@@ -1,7 +1,7 @@
 from fastapi import Depends, HTTPException, Path, status
 from pydantic import UUID4
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.errors import EntityDoesNotExist
 from resources import strings
 from models.domain.shared_service import SharedService

api_app/api/dependencies/workspace_service_templates.py

@@ -1,6 +1,6 @@
 from fastapi import Depends, HTTPException, Path, status
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.errors import EntityDoesNotExist
 from db.repositories.resource_templates import ResourceTemplateRepository
 from models.domain.resource import ResourceType

api_app/api/dependencies/workspaces.py

@@ -1,7 +1,7 @@
 from fastapi import Depends, HTTPException, Path, status
 from pydantic import UUID4
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.errors import EntityDoesNotExist, ResourceIsNotDeployed
 from db.repositories.operations import OperationRepository
 from db.repositories.user_resources import UserResourceRepository

api_app/api/helpers.py

@@ -0,0 +1,22 @@
+from typing import Callable, Type
+
+from fastapi import HTTPException, status
+
+from db.errors import UnableToAccessDatabase
+from db.repositories.base import BaseRepository
+from resources.strings import UNABLE_TO_GET_STATE_STORE_CLIENT
+from services.logging import logger
+
+
+def get_repository(repo_type: Type[BaseRepository],) -> Callable:
+    async def _get_repo() -> BaseRepository:
+        try:
+            return await repo_type.create()
+        except UnableToAccessDatabase:
+            logger.exception(UNABLE_TO_GET_STATE_STORE_CLIENT)
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail=UNABLE_TO_GET_STATE_STORE_CLIENT,
+            )
+
+    return _get_repo

api_app/api/routes/airlock.py

@@ -3,6 +3,7 @@
 from fastapi import APIRouter, Depends, HTTPException, status as status_code, Response
 
 from jsonschema.exceptions import ValidationError
+from api.helpers import get_repository
 from db.repositories.resources_history import ResourceHistoryRepository
 from db.repositories.user_resources import UserResourceRepository
 from db.repositories.workspace_services import WorkspaceServiceRepository
@@ -11,7 +12,6 @@
 from db.repositories.airlock_requests import AirlockRequestRepository
 from db.errors import EntityDoesNotExist, UserNotAuthorizedToUseTemplate
 
-from api.dependencies.database import get_repository
 from api.dependencies.workspaces import get_workspace_by_id_from_path, get_deployed_workspace_by_id_from_path
 from api.dependencies.airlock import get_airlock_request_by_id_from_path
 from models.domain.airlock_request import AirlockRequestStatus, AirlockRequestType

api_app/api/routes/api.py

@@ -5,7 +5,7 @@
 from fastapi.openapi.docs import get_swagger_ui_html, get_swagger_ui_oauth2_redirect_html
 from fastapi.openapi.utils import get_openapi
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.repositories.workspaces import WorkspaceRepository
 from api.routes import health, ping, workspaces, workspace_templates, workspace_service_templates, user_resource_templates, \
     shared_services, shared_service_templates, migrations, costs, airlock, operations, metadata

api_app/api/routes/costs.py

@@ -7,8 +7,8 @@
 from pydantic import UUID4
 
 from models.schemas.costs import get_cost_report_responses, get_workspace_cost_report_responses
-from api.dependencies.database import get_repository
 from core import config
+from api.helpers import get_repository
 from db.repositories.shared_services import SharedServiceRepository
 from db.repositories.user_resources import UserResourceRepository
 from db.repositories.workspace_services import WorkspaceServiceRepository

api_app/api/routes/health.py

@@ -1,5 +1,5 @@
 import asyncio
-from fastapi import APIRouter
+from fastapi import APIRouter, Request
 from core import credentials
 from models.schemas.status import HealthCheck, ServiceStatus, StatusEnum
 from resources import strings
@@ -10,13 +10,13 @@
 
 
 @router.get("/health", name=strings.API_GET_HEALTH_STATUS)
-async def health_check() -> HealthCheck:
+async def health_check(request: Request) -> HealthCheck:
     # The health endpoint checks the status of key components of the system.
     # Note that Resource Processor checks incur Azure management calls, so
     # calling this endpoint frequently may result in API throttling.
-    async with credentials.get_credential_async() as credential:
+    async with credentials.get_credential_async_context() as credential:
         cosmos, sb, rp = await asyncio.gather(
-            create_state_store_status(credential),
+            create_state_store_status(),
             create_service_bus_status(credential),
             create_resource_processor_status(credential)
         )

api_app/api/routes/migrations.py

@@ -1,11 +1,11 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from db.migrations.airlock import AirlockMigration
 from db.migrations.resources import ResourceMigration
+from api.helpers import get_repository
 from db.repositories.operations import OperationRepository
 from db.repositories.resources_history import ResourceHistoryRepository
 from services.authentication import get_current_admin_user
 from resources import strings
-from api.dependencies.database import get_repository
 from db.migrations.shared_services import SharedServiceMigration
 from db.migrations.workspaces import WorkspaceMigration
 from db.repositories.resources import ResourceRepository

api_app/api/routes/operations.py

@@ -1,7 +1,7 @@
 from fastapi import APIRouter, Depends
 
+from api.helpers import get_repository
 from db.repositories.operations import OperationRepository
-from api.dependencies.database import get_repository
 from models.schemas.operation import OperationInList
 from resources import strings
 from services.authentication import get_current_tre_user_or_tre_admin

api_app/api/routes/shared_service_templates.py

@@ -2,7 +2,7 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import parse_obj_as
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.errors import EntityDoesNotExist, EntityVersionExist, InvalidInput
 from db.repositories.resource_templates import ResourceTemplateRepository
 from models.domain.resource import ResourceType

api_app/api/routes/shared_services.py

@@ -5,7 +5,7 @@
 
 from db.repositories.operations import OperationRepository
 from db.errors import DuplicateEntity, MajorVersionUpdateDenied, UserNotAuthorizedToUseTemplate, TargetTemplateVersionDoesNotExist, VersionDowngradeDenied
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from api.dependencies.shared_services import get_shared_service_by_id_from_path, get_operation_by_id_from_path
 from db.repositories.resource_templates import ResourceTemplateRepository
 from db.repositories.resources_history import ResourceHistoryRepository

api_app/api/routes/user_resource_templates.py

@@ -3,10 +3,10 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import parse_obj_as
 
-from api.dependencies.database import get_repository
 from api.dependencies.workspace_service_templates import get_workspace_service_template_by_name_from_path
 from api.routes.resource_helpers import get_template
 from db.errors import EntityVersionExist, InvalidInput
+from api.helpers import get_repository
 from db.repositories.resource_templates import ResourceTemplateRepository
 from models.domain.resource import ResourceType
 from models.schemas.user_resource_template import UserResourceTemplateInResponse, UserResourceTemplateInCreate

api_app/api/routes/workspace_service_templates.py

@@ -2,9 +2,9 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import parse_obj_as
 
-from api.dependencies.database import get_repository
 from api.routes.resource_helpers import get_template
 from db.errors import EntityVersionExist, InvalidInput
+from api.helpers import get_repository
 from db.repositories.resource_templates import ResourceTemplateRepository
 from models.domain.resource import ResourceType
 from models.schemas.resource_template import ResourceTemplateInResponse, ResourceTemplateInformationInList

api_app/api/routes/workspace_templates.py

@@ -2,7 +2,7 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import parse_obj_as
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from db.errors import EntityVersionExist, InvalidInput
 from db.repositories.resource_templates import ResourceTemplateRepository
 from models.domain.resource import ResourceType

api_app/api/routes/workspaces.py

@@ -4,7 +4,7 @@
 
 from jsonschema.exceptions import ValidationError
 
-from api.dependencies.database import get_repository
+from api.helpers import get_repository
 from api.dependencies.workspaces import get_operation_by_id_from_path, get_workspace_by_id_from_path, get_deployed_workspace_by_id_from_path, get_deployed_workspace_service_by_id_from_path, get_workspace_service_by_id_from_path, get_user_resource_by_id_from_path
 from db.errors import InvalidInput, MajorVersionUpdateDenied, TargetTemplateVersionDoesNotExist, UserNotAuthorizedToUseTemplate, VersionDowngradeDenied
 from db.repositories.operations import OperationRepository