From a7e734ab1cfd1051073f1ae506feb03770374faf Mon Sep 17 00:00:00 2001 From: Pankaj Date: Mon, 21 Aug 2023 01:45:26 +0530 Subject: [PATCH 1/5] Add token auth option for azure adx service --- airflow/providers/microsoft/azure/hooks/adx.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/airflow/providers/microsoft/azure/hooks/adx.py b/airflow/providers/microsoft/azure/hooks/adx.py index 83f955753dbf8..2b680df337e48 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 @@ -183,6 +184,14 @@ 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": + credential = conn.password + if not credential: + credential = DefaultAzureCredential() + kcsb = KustoConnectionStringBuilder.with_azure_token_credential( + connection_string=cluster, + credential=credential, + ) else: raise AirflowException(f"Unknown authentication method: {auth_method}") From dcbdecc36253b7b95b9124fe6bfdf9c28d92dec2 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Tue, 22 Aug 2023 02:10:15 +0530 Subject: [PATCH 2/5] clean --- airflow/providers/microsoft/azure/hooks/adx.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/adx.py b/airflow/providers/microsoft/azure/hooks/adx.py index 2b680df337e48..f2cfd1a6c7e53 100644 --- a/airflow/providers/microsoft/azure/hooks/adx.py +++ b/airflow/providers/microsoft/azure/hooks/adx.py @@ -106,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", @@ -185,12 +185,9 @@ 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": - credential = conn.password - if not credential: - credential = DefaultAzureCredential() kcsb = KustoConnectionStringBuilder.with_azure_token_credential( connection_string=cluster, - credential=credential, + credential=DefaultAzureCredential(), ) else: raise AirflowException(f"Unknown authentication method: {auth_method}") From 630f0db948f56c3d7218b05c465200c4c1dbd233 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Thu, 24 Aug 2023 04:15:05 +0530 Subject: [PATCH 3/5] Add test --- .../microsoft/azure/hooks/test_adx.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/providers/microsoft/azure/hooks/test_adx.py b/tests/providers/microsoft/azure/hooks/test_adx.py index ad79d3e0c6484..31f47d135e9ef 100644 --- a/tests/providers/microsoft/azure/hooks/test_adx.py +++ b/tests/providers/microsoft/azure/hooks/test_adx.py @@ -110,6 +110,33 @@ 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 From 8ab50adb4fadd3ec0d2ce4e75d3ed969b0dc7b0a Mon Sep 17 00:00:00 2001 From: Pankaj Date: Thu, 24 Aug 2023 11:21:15 +0530 Subject: [PATCH 4/5] Add test --- tests/providers/microsoft/azure/hooks/test_adx.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/providers/microsoft/azure/hooks/test_adx.py b/tests/providers/microsoft/azure/hooks/test_adx.py index 31f47d135e9ef..9395f479168a4 100644 --- a/tests/providers/microsoft/azure/hooks/test_adx.py +++ b/tests/providers/microsoft/azure/hooks/test_adx.py @@ -126,7 +126,11 @@ def test_conn_method_token_creds(self, mock1): ) with patch.dict( in_dict=os.environ, - values={"AZURE_TENANT_ID": "tenant", "AZURE_CLIENT_ID": "client", "AZURE_CLIENT_SECRET": "secret"} + 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" @@ -134,7 +138,7 @@ def test_conn_method_token_creds(self, mock1): tenant_id="tenant", client_id="client", client_secret="secret", - authority="https://login.microsoftonline.com" + authority="https://login.microsoftonline.com", ) @mock.patch.object(KustoClient, "__init__") From 29e34bc0f071288053dae42a7f021786c4bc0ab2 Mon Sep 17 00:00:00 2001 From: Pankaj Date: Thu, 24 Aug 2023 12:02:56 +0530 Subject: [PATCH 5/5] Add test --- tests/providers/microsoft/azure/hooks/test_adx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/microsoft/azure/hooks/test_adx.py b/tests/providers/microsoft/azure/hooks/test_adx.py index 9395f479168a4..4cf61c440b4ae 100644 --- a/tests/providers/microsoft/azure/hooks/test_adx.py +++ b/tests/providers/microsoft/azure/hooks/test_adx.py @@ -130,7 +130,7 @@ def test_conn_method_token_creds(self, mock1): "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"