diff --git a/airflow/providers/microsoft/azure/hooks/adx.py b/airflow/providers/microsoft/azure/hooks/adx.py index 83f955753dbf8..f2cfd1a6c7e53 100644 --- a/airflow/providers/microsoft/azure/hooks/adx.py +++ b/airflow/providers/microsoft/azure/hooks/adx.py @@ -28,6 +28,7 @@ import warnings from typing import Any +from azure.identity import DefaultAzureCredential from azure.kusto.data import ClientRequestProperties, KustoClient, KustoConnectionStringBuilder from azure.kusto.data.exceptions import KustoServiceError from azure.kusto.data.response import KustoResponseDataSetV2 @@ -105,7 +106,7 @@ def get_ui_field_behaviour() -> dict[str, Any]: "placeholders": { "login": "Varies with authentication method", "password": "Varies with authentication method", - "auth_method": "AAD_APP/AAD_APP_CERT/AAD_CREDS/AAD_DEVICE", + "auth_method": "AAD_APP/AAD_APP_CERT/AAD_CREDS/AAD_DEVICE/AZURE_TOKEN_CRED", "tenant": "Used with AAD_APP/AAD_APP_CERT/AAD_CREDS", "certificate": "Used with AAD_APP_CERT", "thumbprint": "Used with AAD_APP_CERT", @@ -183,6 +184,11 @@ def get_required_param(name: str) -> str: ) elif auth_method == "AAD_DEVICE": kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(cluster) + elif auth_method == "AZURE_TOKEN_CRED": + kcsb = KustoConnectionStringBuilder.with_azure_token_credential( + connection_string=cluster, + credential=DefaultAzureCredential(), + ) else: raise AirflowException(f"Unknown authentication method: {auth_method}") diff --git a/tests/providers/microsoft/azure/hooks/test_adx.py b/tests/providers/microsoft/azure/hooks/test_adx.py index ad79d3e0c6484..4cf61c440b4ae 100644 --- a/tests/providers/microsoft/azure/hooks/test_adx.py +++ b/tests/providers/microsoft/azure/hooks/test_adx.py @@ -110,6 +110,37 @@ def test_conn_method_aad_creds(self, mock_init): ) ) + @mock.patch("azure.identity._credentials.environment.ClientSecretCredential") + def test_conn_method_token_creds(self, mock1): + db.merge_conn( + Connection( + conn_id=ADX_TEST_CONN_ID, + conn_type="azure_data_explorer", + host="https://help.kusto.windows.net", + extra=json.dumps( + { + "auth_method": "AZURE_TOKEN_CRED", + } + ), + ) + ) + with patch.dict( + in_dict=os.environ, + values={ + "AZURE_TENANT_ID": "tenant", + "AZURE_CLIENT_ID": "client", + "AZURE_CLIENT_SECRET": "secret", + }, + ): + hook = AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID) + assert hook.connection._kcsb.data_source == "https://help.kusto.windows.net" + mock1.assert_called_once_with( + tenant_id="tenant", + client_id="client", + client_secret="secret", + authority="https://login.microsoftonline.com", + ) + @mock.patch.object(KustoClient, "__init__") def test_conn_method_aad_app(self, mock_init): mock_init.return_value = None