Fix Azure provider hooks ignoring cloud_environment connection extra (#65320)

Author: sfirke

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)
 

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,13 +82,19 @@ 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.")
             credential = ClientSecretCredential(
                 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:

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"

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,
         )