Repository: microsoft/AzureTRE

Issue:
[Task] API to use Managed Identity to authenticate against Cosmos DB
**Is your feature request related to a problem?

Currently the API use an API key `STATE_STORE_KEY` to authenticate against Cosmos DB. The API uses Managed Identity to authenticate against Azure Service Bus.

**Describe the solution you'd like

Using the same simplified access pattern to authenticate from the API to Azure resource.

Note: looks like Cosmos data-plane doesn't really support managed identity (in Python at least) but there might be a way to do this:
https://github.com/MicrosoftDocs/azure-docs/issues/29762
https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/tutorial-linux-vm-access-cosmos-db



Discussion:
This is the best I came up with:
```python
from azure.cosmos import CosmosClient
from azure.identity import DefaultAzureCredential
from azure.mgmt.cosmosdb import CosmosDBManagementClient

subscription_id = "..."
resource_group = "..."
account_name = "..."

credential = DefaultAzureCredential()
dbmgmt = CosmosDBManagementClient(credential, subscription_id, "https://management.azure.com")
keys = dbmgmt.database_accounts.list_keys(resource_group, account_name)

client = CosmosClient(url="...", credential=keys.primary_master_key)
```

Resolving patch:
diff --git a/management_api_app/api/dependencies/database.py b/management_api_app/api/dependencies/database.py
--- a/management_api_app/api/dependencies/database.py
+++ b/management_api_app/api/dependencies/database.py
@@ -2,6 +2,8 @@
 from typing import Callable, Type
 
 from azure.cosmos import CosmosClient
+from azure.identity import DefaultAzureCredential
+from azure.mgmt.cosmosdb import CosmosDBManagementClient
 from fastapi import Depends, FastAPI, HTTPException
 from starlette.requests import Request
 from starlette.status import HTTP_503_SERVICE_UNAVAILABLE
@@ -16,18 +18,30 @@ def connect_to_db() -> CosmosClient:
     logging.debug(f"Connecting to {config.STATE_STORE_ENDPOINT}")
 
     try:
+        primary_master_key = get_store_key()
         if config.DEBUG:
             # ignore TLS(setup is pain) when on dev container and connecting to cosmosdb on windows host.
-            cosmos_client = CosmosClient(config.STATE_STORE_ENDPOINT, config.STATE_STORE_KEY,
-                                         connection_verify=False)
+            cosmos_client = CosmosClient(config.STATE_STORE_ENDPOINT, primary_master_key, connection_verify=False)
         else:
-            cosmos_client = CosmosClient(config.STATE_STORE_ENDPOINT, config.STATE_STORE_KEY)
+            cosmos_client = CosmosClient(config.STATE_STORE_ENDPOINT, primary_master_key)
         logging.debug("Connection established")
         return cosmos_client
     except Exception as e:
         logging.debug(f"Connection to state store could not be established: {e}")
 
 
+def get_store_key() -> str:
+    if config.STATE_STORE_KEY:
+        primary_master_key = config.STATE_STORE_KEY
+    else:
+        credential = DefaultAzureCredential(managed_identity_client_id=config.MANAGED_IDENTITY_CLIENT_ID, exclude_shared_token_cache_credential=True)
+        cosmosdb_client = CosmosDBManagementClient(credential, subscription_id=config.SUBSCRIPTION_ID)
+        database_keys = cosmosdb_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
+
+
 def get_db_client(app: FastAPI) -> CosmosClient:
     if not app.state.cosmos_client:
         app.state.cosmos_client = connect_to_db()
diff --git a/management_api_app/core/config.py b/management_api_app/core/config.py
--- a/management_api_app/core/config.py
+++ b/management_api_app/core/config.py
@@ -16,9 +16,13 @@
 # State store configuration
 STATE_STORE_ENDPOINT: str = config("STATE_STORE_ENDPOINT", default="")      # Cosmos DB endpoint
 STATE_STORE_KEY: str = config("STATE_STORE_KEY", default="")                # Cosmos DB access key
+COSMOSDB_ACCOUNT_NAME: str = config("COSMOSDB_ACCOUNT_NAME", default="")                # Cosmos DB account name
 STATE_STORE_DATABASE = "AzureTRE"
 STATE_STORE_RESOURCES_CONTAINER = "Resources"
 STATE_STORE_RESOURCE_TEMPLATES_CONTAINER = "ResourceTemplates"
+SUBSCRIPTION_ID: str = config("SUBSCRIPTION_ID", default="")
+RESOURCE_GROUP_NAME: str = config("RESOURCE_GROUP_NAME", default="")
+
 
 # Service bus configuration
 SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE: str = config("SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE", default="")
diff --git a/management_api_app/services/health_checker.py b/management_api_app/services/health_checker.py
--- a/management_api_app/services/health_checker.py
+++ b/management_api_app/services/health_checker.py
@@ -1,7 +1,8 @@
 from azure.core import exceptions
 from azure.cosmos import CosmosClient
 
-from core.config import STATE_STORE_ENDPOINT, STATE_STORE_KEY
+from api.dependencies.database import get_store_key
+from core import config
 from models.schemas.status import StatusEnum
 from resources import strings
 
@@ -10,7 +11,8 @@ def create_state_store_status() -> (StatusEnum, str):
     status = StatusEnum.ok
     message = ""
     try:
-        client = CosmosClient(STATE_STORE_ENDPOINT, STATE_STORE_KEY)    # noqa: F841 - flake 8 client is not used
+        primary_master_key = get_store_key()
+        client = CosmosClient(config.STATE_STORE_ENDPOINT, primary_master_key)    # noqa: F841 - flake 8 client is not used
     except exceptions.ServiceRequestError:
         status = StatusEnum.not_ok
         message = strings.STATE_STORE_ENDPOINT_NOT_RESPONDING


Test patch:
diff --git a/management_api_app/tests_ma/test_services/test_health_checker.py b/management_api_app/tests_ma/test_services/test_health_checker.py
--- a/management_api_app/tests_ma/test_services/test_health_checker.py
+++ b/management_api_app/tests_ma/test_services/test_health_checker.py
@@ -7,7 +7,9 @@
 
 
 @patch("services.health_checker.CosmosClient")
-def test_get_state_store_status_responding(cosmos_client_mock) -> None:
+@patch("services.health_checker.get_store_key")
+def test_get_state_store_status_responding(cosmos_client_mock, get_store_key_mock) -> None:
+    get_store_key_mock.return_value = None
     cosmos_client_mock.return_value = None
 
     status, message = health_checker.create_state_store_status()
@@ -17,7 +19,9 @@ def test_get_state_store_status_responding(cosmos_client_mock) -> None:
 
 
 @patch("services.health_checker.CosmosClient")
-def test_get_state_store_status_not_responding(cosmos_client_mock) -> None:
+@patch("services.health_checker.get_store_key")
+def test_get_state_store_status_not_responding(cosmos_client_mock, get_store_key_mock) -> None:
+    get_store_key_mock.return_value = None
     cosmos_client_mock.return_value = None
     cosmos_client_mock.side_effect = ServiceRequestError(message="some message")
 
@@ -28,7 +32,9 @@ def test_get_state_store_status_not_responding(cosmos_client_mock) -> None:
 
 
 @patch("services.health_checker.CosmosClient")
-def test_get_state_store_status_other_exception(cosmos_client_mock) -> None:
+@patch("services.health_checker.get_store_key")
+def test_get_state_store_status_other_exception(cosmos_client_mock, get_store_key_mock) -> None:
+    get_store_key_mock.return_value = None
     cosmos_client_mock.return_value = None
     cosmos_client_mock.side_effect = Exception()