You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2023/11/01 08:40:35 UTC
(airflow) branch main updated: feat(provider/azure): add managed identity support to cosmos hook (#35323)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a13c091f2d feat(provider/azure): add managed identity support to cosmos hook (#35323)
a13c091f2d is described below
commit a13c091f2d7177e54de023ae3ab3e4b3fda4d695
Author: Wei Lee <we...@gmail.com>
AuthorDate: Wed Nov 1 16:40:28 2023 +0800
feat(provider/azure): add managed identity support to cosmos hook (#35323)
---
airflow/providers/microsoft/azure/hooks/cosmos.py | 20 ++++++++++++++++----
.../connections/azure_cosmos.rst | 11 +++++++++--
2 files changed, 25 insertions(+), 6 deletions(-)
diff --git a/airflow/providers/microsoft/azure/hooks/cosmos.py b/airflow/providers/microsoft/azure/hooks/cosmos.py
index 1f9883899e..596faf6b04 100644
--- a/airflow/providers/microsoft/azure/hooks/cosmos.py
+++ b/airflow/providers/microsoft/azure/hooks/cosmos.py
@@ -31,12 +31,11 @@ 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, AirflowException
from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import get_field
+from airflow.providers.microsoft.azure.utils import get_default_azure_credential, get_field
class AzureCosmosDBHook(BaseHook):
@@ -78,6 +77,12 @@ class AzureCosmosDBHook(BaseHook):
lazy_gettext("Resource Group Name (optional)"),
widget=BS3TextFieldWidget(),
),
+ "managed_identity_client_id": StringField(
+ lazy_gettext("Managed Identity Client ID"), widget=BS3TextFieldWidget()
+ ),
+ "workload_identity_tenant_id": StringField(
+ lazy_gettext("Workload Identity Tenant ID"), widget=BS3TextFieldWidget()
+ ),
}
@staticmethod
@@ -96,6 +101,8 @@ class AzureCosmosDBHook(BaseHook):
"collection_name": "collection name",
"subscription_id": "Subscription ID (required for Azure AD authentication)",
"resource_group_name": "Resource Group Name (required for Azure AD authentication)",
+ "managed_identity_client_id": "Managed Identity Client ID",
+ "workload_identity_tenant_id": "Workload Identity Tenant ID",
},
}
@@ -126,9 +133,14 @@ class AzureCosmosDBHook(BaseHook):
if conn.password:
master_key = conn.password
elif resource_group_name:
+ managed_identity_client_id = self._get_field(extras, "managed_identity_client_id")
+ workload_identity_tenant_id = self._get_field(extras, "workload_identity_tenant_id")
+ subscritption_id = self._get_field(extras, "subscription_id")
management_client = CosmosDBManagementClient(
- credential=DefaultAzureCredential(),
- subscription_id=self._get_field(extras, "subscription_id"),
+ credential=get_default_azure_credential(
+ managed_identity_client_id, workload_identity_tenant_id
+ ),
+ subscription_id=subscritption_id,
)
database_account = urlparse(conn.login).netloc.split(".")[0]
diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst b/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst
index 3c54b98bc2..e7933fc897 100644
--- a/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst
+++ b/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst
@@ -27,11 +27,12 @@ The Microsoft Azure Cosmos connection type enables the Azure Cosmos Integrations
Authenticating to Azure
-----------------------
-There are two ways to connect to Azure Cosmos using Airflow.
+There are three ways to connect to Azure Cosmos using Airflow.
1. Use `Primary Keys`_
i.e. add specific credentials (client_id, secret, tenant) and account name to the Airflow connection.
-2. Fallback on DefaultAzureCredential_.
+2. Use managed identity by setting ``managed_identity_client_id``, ``workload_identity_tenant_id`` (under the hook, it uses DefaultAzureCredential_ with these arguments)
+3. Fallback on DefaultAzureCredential_.
This includes a mechanism to try different options to authenticate: Managed System Identity, environment variables, authentication through Azure CLI, etc.
Default Connection IDs
@@ -57,6 +58,8 @@ Extra (optional)
* ``collection_name``: Specify the azure cosmos collection to use.
* ``subscription_id``: Specify the ID of the subscription used for the initial connection. Required for falling back to DefaultAzureCredential_
* ``resource_group_name``: Specify the Azure Resource Group Name under which the desired azure cosmos resides. Required for falling back to DefaultAzureCredential_
+ * ``managed_identity_client_id``: The client ID of a user-assigned managed identity. If provided with `workload_identity_tenant_id`, they'll pass to ``DefaultAzureCredential``.
+ * ``workload_identity_tenant_id``: ID of the application's Microsoft Entra tenant. Also called its "directory" ID. If provided with `managed_identity_client_id`, they'll pass to ``DefaultAzureCredential``.
When specifying the connection in environment variable you should specify
@@ -73,3 +76,7 @@ For example:
.. _Primary Keys: https://docs.microsoft.com/en-us/azure/cosmos-db/secure-access-to-data#primary-keys
.. _DefaultAzureCredential: https://docs.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential
+
+.. spelling:word-list::
+
+ Entra