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 9dac4878c7 Add DefaultAzureCredential auth for ADX service (#33627)
9dac4878c7 is described below
commit 9dac4878c70f4178b89c2f7667b0d8ca0ca7dff7
Author: Pankaj Singh <[email protected]>
AuthorDate: Thu Aug 24 14:41:52 2023 +0530
Add DefaultAzureCredential auth for ADX service (#33627)
---
airflow/providers/microsoft/azure/hooks/adx.py | 8 +++++-
tests/providers/microsoft/azure/hooks/test_adx.py | 31 +++++++++++++++++++++++
2 files changed, 38 insertions(+), 1 deletion(-)
diff --git a/airflow/providers/microsoft/azure/hooks/adx.py
b/airflow/providers/microsoft/azure/hooks/adx.py
index 83f955753d..f2cfd1a6c7 100644
--- a/airflow/providers/microsoft/azure/hooks/adx.py
+++ b/airflow/providers/microsoft/azure/hooks/adx.py
@@ -28,6 +28,7 @@ from __future__ import annotations
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 @@ class AzureDataExplorerHook(BaseHook):
"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 @@ class AzureDataExplorerHook(BaseHook):
)
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 ad79d3e0c6..4cf61c440b 100644
--- a/tests/providers/microsoft/azure/hooks/test_adx.py
+++ b/tests/providers/microsoft/azure/hooks/test_adx.py
@@ -110,6 +110,37 @@ class TestAzureDataExplorerHook:
)
)
+
@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
