diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py index 7bbe931ee4858..6ed931f6dd6d2 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/base_azure.py @@ -20,7 +20,7 @@ from azure.common.client_factory import get_client_from_auth_file, get_client_from_json_dict from azure.common.credentials import ServicePrincipalCredentials -from azure.identity import ClientSecretCredential, DefaultAzureCredential +from azure.identity import AzureAuthorityHosts, ClientSecretCredential, DefaultAzureCredential from airflow.providers.common.compat.sdk import AirflowException, BaseHook from airflow.providers.microsoft.azure.utils import ( @@ -34,6 +34,24 @@ from airflow.sdk import Connection +_AZURE_CLOUD_ENVIRONMENTS: dict[str, dict[str, Any]] = { + "AzurePublicCloud": { + "authority": AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, + "base_url": "https://management.azure.com", + "credential_scopes": ["https://management.azure.com/.default"], + }, + "AzureUSGovernment": { + "authority": AzureAuthorityHosts.AZURE_GOVERNMENT, + "base_url": "https://management.usgovcloudapi.net", + "credential_scopes": ["https://management.usgovcloudapi.net/.default"], + }, + "AzureChinaCloud": { + "authority": AzureAuthorityHosts.AZURE_CHINA, + "base_url": "https://management.chinacloudapi.cn", + "credential_scopes": ["https://management.chinacloudapi.cn/.default"], + }, +} + class AzureBaseHook(BaseHook): """ @@ -63,6 +81,9 @@ def get_connection_form_widgets(cls) -> dict[str, Any]: return { "tenantId": StringField(lazy_gettext("Azure Tenant ID"), widget=BS3TextFieldWidget()), "subscriptionId": StringField(lazy_gettext("Azure Subscription ID"), widget=BS3TextFieldWidget()), + "cloud_environment": StringField( + lazy_gettext("Azure Cloud Environment"), widget=BS3TextFieldWidget() + ), } @classmethod @@ -88,6 +109,7 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: "password": "secret (token credentials auth)", "tenantId": "tenantId (token credentials auth)", "subscriptionId": "subscriptionId (token credentials auth)", + "cloud_environment": "AzurePublicCloud (default) | AzureUSGovernment | AzureChinaCloud", }, } @@ -163,11 +185,18 @@ def _get_client_secret_credential( extra_dejson = conn.extra_dejson tenant = extra_dejson.get("tenantId") use_azure_identity_object = extra_dejson.get("use_azure_identity_object", False) + + cloud_env_name = extra_dejson.get("cloud_environment", "AzurePublicCloud") + cloud_env = _AZURE_CLOUD_ENVIRONMENTS.get( + cloud_env_name, _AZURE_CLOUD_ENVIRONMENTS["AzurePublicCloud"] + ) + if use_azure_identity_object: return ClientSecretCredential( client_id=conn.login, # type: ignore[arg-type] client_secret=conn.password, # type: ignore[arg-type] tenant_id=tenant, # type: ignore[arg-type] + authority=cloud_env["authority"], ) return ServicePrincipalCredentials(client_id=conn.login, secret=conn.password, tenant=tenant) diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py index b30d729c47881..27941ec53afc8 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py @@ -26,7 +26,7 @@ from azure.mgmt.containerinstance import ContainerInstanceManagementClient from airflow.providers.common.compat.sdk import AirflowException -from airflow.providers.microsoft.azure.hooks.base_azure import AzureBaseHook +from airflow.providers.microsoft.azure.hooks.base_azure import _AZURE_CLOUD_ENVIRONMENTS, AzureBaseHook from airflow.providers.microsoft.azure.utils import get_sync_default_azure_credential if TYPE_CHECKING: @@ -82,6 +82,11 @@ def get_conn(self) -> Any: self.log.info("Getting connection using a JSON config.") return get_client_from_json_dict(client_class=self.sdk_client, config_dict=key_json) + cloud_env_name = conn.extra_dejson.get("cloud_environment", "AzurePublicCloud") + cloud_env = _AZURE_CLOUD_ENVIRONMENTS.get( + cloud_env_name, _AZURE_CLOUD_ENVIRONMENTS["AzurePublicCloud"] + ) + credential: ClientSecretCredential | DefaultAzureCredential if all([conn.login, conn.password, tenant]): self.log.info("Getting connection using specific credentials and subscription_id.") @@ -89,6 +94,7 @@ def get_conn(self) -> Any: client_id=cast("str", conn.login), client_secret=cast("str", conn.password), tenant_id=cast("str", tenant), + authority=cloud_env["authority"], ) else: self.log.info("Using DefaultAzureCredential as credential") @@ -103,6 +109,8 @@ def get_conn(self) -> Any: return ContainerInstanceManagementClient( credential=credential, subscription_id=subscription_id, + base_url=cloud_env["base_url"], + credential_scopes=cloud_env["credential_scopes"], ) def create_or_update(self, resource_group: str, name: str, container_group: ContainerGroup) -> None: diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_base_azure.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_base_azure.py index a0fc50857d984..3e4a7fc27b8f5 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_base_azure.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_base_azure.py @@ -20,6 +20,7 @@ from unittest.mock import MagicMock, Mock, patch import pytest +from azure.identity import AzureAuthorityHosts from airflow.models import Connection from airflow.providers.microsoft.azure.hooks.base_azure import AzureBaseHook @@ -156,6 +157,7 @@ def test_get_credential_with_client_secret(self, mock_spc, mocked_connection): client_id=mocked_connection.login, client_secret=mocked_connection.password, tenant_id=mocked_connection.extra_dejson["tenantId"], + authority=AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, ) assert cred == "foo-bar" @@ -221,3 +223,41 @@ def test_get_token_with_azure_default_credential(self, mock_spc, mocked_connecti mock_spc.assert_called_once_with() assert token == "new-token" + + @patch(f"{MODULE}.ClientSecretCredential") + @pytest.mark.parametrize( + ("cloud_env", "expected_authority"), + [ + pytest.param(None, AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, id="default_public_cloud"), + pytest.param( + "AzurePublicCloud", AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, id="explicit_public_cloud" + ), + pytest.param("AzureUSGovernment", AzureAuthorityHosts.AZURE_GOVERNMENT, id="us_government"), + pytest.param("AzureChinaCloud", AzureAuthorityHosts.AZURE_CHINA, id="china_cloud"), + ], + ) + def test_get_credential_cloud_environment( + self, mock_csc, cloud_env, expected_authority, create_mock_connection + ): + extras = {"tenantId": "my_tenant", "use_azure_identity_object": True} + if cloud_env is not None: + extras["cloud_environment"] = cloud_env + + create_mock_connection( + Connection( + conn_id="azure_default", + login="my_login", + password="my_password", + extra=extras, + ) + ) + mock_csc.return_value = "credential" + cred = AzureBaseHook().get_credential() + + mock_csc.assert_called_once_with( + client_id="my_login", + client_secret="my_password", + tenant_id="my_tenant", + authority=expected_authority, + ) + assert cred == "credential" diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py index a9510e9094501..0bebf4d7afc59 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py @@ -21,6 +21,7 @@ import pytest from azure.core.exceptions import ResourceNotFoundError +from azure.identity import AzureAuthorityHosts from azure.mgmt.containerinstance.models import ( Logs, ResourceRequests, @@ -165,4 +166,89 @@ def test_get_conn_fallback_to_default_azure_credential( mock_client_cls.assert_called_once_with( credential=mock_credential, subscription_id="subscription_id", + base_url="https://management.azure.com", + credential_scopes=["https://management.azure.com/.default"], + ) + + +class TestAzureContainerInstanceHookCloudEnvironment: + @pytest.mark.parametrize( + ("cloud_env", "expected_authority", "expected_base_url", "expected_scopes"), + [ + pytest.param( + None, + AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, + "https://management.azure.com", + ["https://management.azure.com/.default"], + id="default_public_cloud", + ), + pytest.param( + "AzurePublicCloud", + AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, + "https://management.azure.com", + ["https://management.azure.com/.default"], + id="explicit_public_cloud", + ), + pytest.param( + "AzureUSGovernment", + AzureAuthorityHosts.AZURE_GOVERNMENT, + "https://management.usgovcloudapi.net", + ["https://management.usgovcloudapi.net/.default"], + id="us_government", + ), + pytest.param( + "AzureChinaCloud", + AzureAuthorityHosts.AZURE_CHINA, + "https://management.chinacloudapi.cn", + ["https://management.chinacloudapi.cn/.default"], + id="china_cloud", + ), + ], + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.ContainerInstanceManagementClient") + @patch("airflow.providers.microsoft.azure.hooks.container_instance.ClientSecretCredential") + def test_get_conn_cloud_environment( + self, + mock_credential_cls, + mock_client_cls, + cloud_env, + expected_authority, + expected_base_url, + expected_scopes, + create_mock_connection, + ): + extras = { + "tenantId": "my-tenant", + "subscriptionId": "my-subscription", + } + if cloud_env is not None: + extras["cloud_environment"] = cloud_env + + mock_connection = create_mock_connection( + Connection( + conn_id="azure_container_instance_cloud_test", + conn_type="azure_container_instances", + login="my-client-id", + password="my-secret", + extra=extras, + ) + ) + + mock_credential_cls.return_value = MagicMock() + mock_client_cls.return_value = MagicMock() + + hook = AzureContainerInstanceHook(azure_conn_id=mock_connection.conn_id) + hook.get_conn() + + mock_credential_cls.assert_called_once_with( + client_id="my-client-id", + client_secret="my-secret", + tenant_id="my-tenant", + authority=expected_authority, + ) + mock_client_cls.assert_called_once_with( + credential=mock_credential_cls.return_value, + subscription_id="my-subscription", + base_url=expected_base_url, + credential_scopes=expected_scopes, )