pankajastro commented on code in PR #73374:
URL: https://github.com/apache/airflow/pull/73374#discussion_r4091706885


##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -153,6 +154,14 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, 
Any]:
                     conf[key] = conn.extra_dejson[key]
             return conf
 
+        def _get_wasb_extra_field(extra_dejson: dict[str, Any], field_name: 
str) -> Any:
+            # Older Airflow connection UIs wrote custom extra fields as
+            # extra__wasb__<field_name> instead of the bare key; WasbHook 
still reads that
+            # legacy spelling as a fallback, so this must too.
+            if field_name in extra_dejson:
+                return extra_dejson[field_name]
+            return extra_dejson.get(f"extra__wasb__{field_name}")

Review Comment:
   Done, thanks — cleaner.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -196,6 +205,82 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, 
Any]:
                     key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
                 credentials = self._remove_none_values({"key_path": key_path, 
"keyfile_dict": keyfile_dict})
 
+            case "wasb":
+                extra_dejson = conn.extra_dejson
+                for unsupported_field in (
+                    "connection_string",
+                    "managed_identity_client_id",
+                    "workload_identity_tenant_id",
+                ):
+                    if _get_wasb_extra_field(extra_dejson, unsupported_field):
+                        raise ValueError(
+                            f"Connection field {unsupported_field!r} is not 
supported for DataFusion "
+                            "Azure Blob Storage access; only 
tenant_id+login+password (service "
+                            "principal), sas_token, 
shared_access_key/account_key/password, or ambient "
+                            "credentials (AZURE_* environment variables, 
managed identity, workload "
+                            "identity, or az login) are used."
+                        )
+                credentials = {"account": 
self._resolve_wasb_account(conn.host, conn.login)}
+                tenant_id = _get_wasb_extra_field(extra_dejson, "tenant_id")
+                sas_token = _get_wasb_extra_field(extra_dejson, "sas_token")
+                explicit_credential = False
+                if tenant_id:

Review Comment:
   Done, thanks.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -196,6 +205,82 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, 
Any]:
                     key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
                 credentials = self._remove_none_values({"key_path": key_path, 
"keyfile_dict": keyfile_dict})
 
+            case "wasb":
+                extra_dejson = conn.extra_dejson
+                for unsupported_field in (
+                    "connection_string",
+                    "managed_identity_client_id",
+                    "workload_identity_tenant_id",
+                ):
+                    if _get_wasb_extra_field(extra_dejson, unsupported_field):
+                        raise ValueError(
+                            f"Connection field {unsupported_field!r} is not 
supported for DataFusion "
+                            "Azure Blob Storage access; only 
tenant_id+login+password (service "
+                            "principal), sas_token, 
shared_access_key/account_key/password, or ambient "
+                            "credentials (AZURE_* environment variables, 
managed identity, workload "
+                            "identity, or az login) are used."
+                        )
+                credentials = {"account": 
self._resolve_wasb_account(conn.host, conn.login)}
+                tenant_id = _get_wasb_extra_field(extra_dejson, "tenant_id")
+                sas_token = _get_wasb_extra_field(extra_dejson, "sas_token")
+                explicit_credential = False
+                if tenant_id:
+                    if not conn.login or not conn.password:
+                        # Falling through here would silently switch identity 
(ambient auth, or
+                        # the client secret sent as a shared key) instead of 
failing clearly.
+                        missing = "login (client_id)" if not conn.login else 
"password (client_secret)"
+                        raise ValueError(
+                            f"Connection extra 'tenant_id' is set for 
DataFusion Azure Blob Storage "
+                            f"service-principal auth, but {missing} is not."
+                        )
+                    credentials.update(
+                        {"client_id": conn.login, "client_secret": 
conn.password, "tenant_id": tenant_id}
+                    )
+                    explicit_credential = True
+                elif sas_token:

Review Comment:
   Done, thanks.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -196,6 +205,82 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, 
Any]:
                     key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
                 credentials = self._remove_none_values({"key_path": key_path, 
"keyfile_dict": keyfile_dict})
 
+            case "wasb":
+                extra_dejson = conn.extra_dejson
+                for unsupported_field in (
+                    "connection_string",
+                    "managed_identity_client_id",
+                    "workload_identity_tenant_id",
+                ):
+                    if _get_wasb_extra_field(extra_dejson, unsupported_field):
+                        raise ValueError(
+                            f"Connection field {unsupported_field!r} is not 
supported for DataFusion "
+                            "Azure Blob Storage access; only 
tenant_id+login+password (service "
+                            "principal), sas_token, 
shared_access_key/account_key/password, or ambient "
+                            "credentials (AZURE_* environment variables, 
managed identity, workload "
+                            "identity, or az login) are used."
+                        )
+                credentials = {"account": 
self._resolve_wasb_account(conn.host, conn.login)}
+                tenant_id = _get_wasb_extra_field(extra_dejson, "tenant_id")
+                sas_token = _get_wasb_extra_field(extra_dejson, "sas_token")
+                explicit_credential = False

Review Comment:
   Not quite — the `else` branch's inner `if access_key:` can be false too 
(nothing set at all, i.e. pure ambient auth), which is exactly the case the 
env-precedence guard below needs to skip. 
`test_get_credentials_azure_env_precedence_guard_ignores_pure_ambient_auth` 
pins this: env var set + no credential at all must **not** raise. Removing the 
flag would make that guard fire on ambient-only connections too, so keeping it.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -107,6 +107,37 @@ def get_scheme(self) -> str:
         return "gs://"
 
 
+class AzureObjectStorageProvider(ObjectStorageProvider):
+    """Azure Object Storage Provider using DataFusion's MicrosoftAzure."""
+
+    @property
+    def get_storage_type(self) -> StorageType:

Review Comment:
   This mirrors `S3ObjectStorageProvider`/`GCSObjectStorageProvider` exactly, 
and the abstract base (`ObjectStorageProvider.get_storage_type`) itself 
declares it as a `@property` — Azure can't unilaterally drop it without 
breaking the shared interface. Happy to revisit as a separate repo-wide cleanup 
touching the base class and all four providers, but out of scope here.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -107,6 +107,37 @@ def get_scheme(self) -> str:
         return "gs://"
 
 
+class AzureObjectStorageProvider(ObjectStorageProvider):
+    """Azure Object Storage Provider using DataFusion's MicrosoftAzure."""
+
+    @property
+    def get_storage_type(self) -> StorageType:
+        """Return the storage type."""

Review Comment:
   Same as the sibling comment on the property — matches the existing S3/GCS 
docstring convention (`"""Return the storage type."""`); would want to drop it 
everywhere at once if we do, not just here.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to