diff --git a/airflow/providers/microsoft/azure/hooks/cosmos.py b/airflow/providers/microsoft/azure/hooks/cosmos.py index 3a73eeada7d7d..2ebf9ed0f9bee 100644 --- a/airflow/providers/microsoft/azure/hooks/cosmos.py +++ b/airflow/providers/microsoft/azure/hooks/cosmos.py @@ -28,11 +28,14 @@ import json import uuid from typing import Any +from urllib.parse import urlparse from azure.cosmos.cosmos_client import CosmosClient from azure.cosmos.exceptions import CosmosHttpResponseError +from azure.identity import DefaultAzureCredential +from azure.mgmt.cosmosdb import CosmosDBManagementClient -from airflow.exceptions import AirflowBadRequest +from airflow.exceptions import AirflowBadRequest, AirflowException from airflow.hooks.base import BaseHook from airflow.providers.microsoft.azure.utils import get_field @@ -68,6 +71,14 @@ def get_connection_form_widgets() -> dict[str, Any]: "collection_name": StringField( lazy_gettext("Cosmos Collection Name (optional)"), widget=BS3TextFieldWidget() ), + "subscription_id": StringField( + lazy_gettext("Subscription ID (optional)"), + widget=BS3TextFieldWidget(), + ), + "resource_group_name": StringField( + lazy_gettext("Resource Group Name (optional)"), + widget=BS3TextFieldWidget(), + ), } @staticmethod @@ -81,9 +92,11 @@ def get_ui_field_behaviour() -> dict[str, Any]: }, "placeholders": { "login": "endpoint uri", - "password": "master key", + "password": "master key (not needed for Azure AD authentication)", "database_name": "database name", "collection_name": "collection name", + "subscription_id": "Subscription ID (required for Azure AD authentication)", + "resource_group_name": "Resource Group Name (required for Azure AD authentication)", }, } @@ -109,7 +122,23 @@ def get_conn(self) -> CosmosClient: conn = self.get_connection(self.conn_id) extras = conn.extra_dejson endpoint_uri = conn.login - master_key = conn.password + resource_group_name = self._get_field(extras, "resource_group_name") + + if conn.password: + master_key = conn.password + elif resource_group_name: + management_client = CosmosDBManagementClient( + credential=DefaultAzureCredential(), + subscription_id=self._get_field(extras, "subscription_id"), + ) + + database_account = urlparse(conn.login).netloc.split(".")[0] + database_account_keys = management_client.database_accounts.list_keys( + resource_group_name, database_account + ) + master_key = database_account_keys.primary_master_key + else: + raise AirflowException("Either password or resource_group_name is required") self.default_database_name = self._get_field(extras, "database_name") self.default_collection_name = self._get_field(extras, "collection_name") diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index a0ed5fbdf2599..8459d583ee9a1 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -66,6 +66,7 @@ dependencies: - apache-airflow>=2.4.0 - azure-batch>=8.0.0 - azure-cosmos>=4.0.0 + - azure-mgmt-cosmosdb - azure-datalake-store>=0.0.45 - azure-identity>=1.3.1 - azure-keyvault-secrets>=4.1.0 diff --git a/docs/apache-airflow-providers-microsoft-azure/index.rst b/docs/apache-airflow-providers-microsoft-azure/index.rst index 73f9a0733c656..a3889484100ff 100644 --- a/docs/apache-airflow-providers-microsoft-azure/index.rst +++ b/docs/apache-airflow-providers-microsoft-azure/index.rst @@ -107,6 +107,7 @@ PIP package Version required ``apache-airflow`` ``>=2.4.0`` ``azure-batch`` ``>=8.0.0`` ``azure-cosmos`` ``>=4.0.0`` +``azure-mgmt-cosmosdb`` ``azure-datalake-store`` ``>=0.0.45`` ``azure-identity`` ``>=1.3.1`` ``azure-keyvault-secrets`` ``>=4.1.0`` diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 5191600a98392..b0dcab34c1747 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -560,6 +560,7 @@ "azure-keyvault-secrets>=4.1.0", "azure-kusto-data>=4.1.0", "azure-mgmt-containerinstance>=1.5.0,<2.0", + "azure-mgmt-cosmosdb", "azure-mgmt-datafactory>=1.0.0,<2.0", "azure-mgmt-datalake-store>=0.5.0", "azure-mgmt-resource>=2.2.0", diff --git a/tests/providers/microsoft/azure/hooks/test_azure_cosmos.py b/tests/providers/microsoft/azure/hooks/test_azure_cosmos.py index af649f1d6e4d9..787ea5c4537c8 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_cosmos.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_cosmos.py @@ -34,7 +34,6 @@ class TestAzureCosmosDbHook: - # Set up an environment to test with def setup_method(self): # set up some test variables @@ -266,6 +265,8 @@ def test_get_ui_field_behaviour_placeholders(self): "password", "database_name", "collection_name", + "subscription_id", + "resource_group_name", ] if get_provider_min_airflow_version("apache-airflow-providers-microsoft-azure") >= (2, 5): raise Exception(