You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "pankajastro (via GitHub)" <gi...@apache.org> on 2026/09/19 05:20:07 UTC

[PR] Add GCS support to the DataFusion object storage layer [airflow]

pankajastro opened a new pull request, #73370:
URL: https://github.com/apache/airflow/pull/73370

   Part of AFL-214's second defect (DataFusion credential resolution only handled AWS, so GCS/Azure got nothing even though DataFusion's storage layer supports them). This PR covers GCS; Azure is left for a separate follow-up.
   
   `DataFusionEngine._get_credentials` only had a `case "aws":` arm, so a GCS-backed `DataSourceConfig` failed with `Unknown connection type google_cloud_platform` at registration, even though nothing prevented constructing one (`gs://` was never recognized by `_extract_storage_type` either).
   
   Adds:
   - `StorageType.GCS` and `gs://` URI recognition in `DataSourceConfig`
   - `GCSObjectStorageProvider` using DataFusion's `GoogleCloud` object store binding
   - A `google_cloud_platform` branch in `_get_credentials` resolving the connection's `key_path` field, falling back to Application Default Credentials when unset (DataFusion's `GoogleCloud` binding only accepts a `service_account_path`, not inline JSON credentials, so `keyfile_dict`-based connections aren't covered — ADC is the fallback for those)
   - `google` extra in `pyproject.toml`, docs updated to mention GCS alongside S3
   
   ---
   
   ##### Was generative AI tooling used to co-author this PR?
   
   - [X] Yes — Claude Code (Sonnet 5)
   
   Generated-by: Claude Code (Sonnet 5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)


-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Add GCS support to the DataFusion object storage layer [airflow]

Posted by "pankajastro (via GitHub)" <gi...@apache.org>.
pankajastro commented on PR #73370:
URL: https://github.com/apache/airflow/pull/73370#issuecomment-5750528430

   Thanks! Updated the description to spell out the actual fallback chain (key_path/keyfile_dict → GOOGLE_APPLICATION_CREDENTIALS → gcloud ADC/metadata server). On the docs note — operators.rst already lists GCS as supported in this same merged commit, so that should be resolved on main


-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Add GCS support to the DataFusion object storage layer [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #73370:
URL: https://github.com/apache/airflow/pull/73370#discussion_r4054712850


##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -53,6 +57,53 @@ def get_scheme(self) -> str:
         return "s3://"
 
 
+class GCSObjectStorageProvider(ObjectStorageProvider):
+    """GCS Object Storage Provider using DataFusion's GoogleCloud."""
+
+    @property
+    def get_storage_type(self) -> StorageType:
+        """Return the storage type."""
+        return StorageType.GCS
+
+    def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None):
+        """Create a GCS object store using DataFusion's GoogleCloud."""
+        if connection_config is None:
+            raise ValueError(f"connection_config must be provided for {self.get_storage_type}")
+
+        credentials = connection_config.credentials
+        key_path = credentials.get("key_path")
+        keyfile_dict = credentials.get("keyfile_dict")
+        temp_key_path: str | None = None
+
+        try:
+            bucket = self.get_bucket(path)
+
+            if not key_path and keyfile_dict:
+                # DataFusion's GoogleCloud binding only accepts a file path, not inline JSON,
+                # so materialize keyfile_dict to a temp file. The credentials file is read once
+                # at construction time and never touched again, so it's safe to delete right after.
+                key_content = keyfile_dict if isinstance(keyfile_dict, str) else json.dumps(keyfile_dict)
+                with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as key_file:
+                    key_file.write(key_content)
+                temp_key_path = key_path = key_file.name
+
+            gcs_store = GoogleCloud(bucket_name=bucket, service_account_path=key_path)

Review Comment:
   The DataFusion binding wraps `build()` in `.expect(...)` ([store.rs:164](https://github.com/apache/datafusion-python/blob/51.0.0/src/store.rs#L164)), so a missing or malformed key file surfaces as `pyo3_runtime.PanicException`, which subclasses `BaseException`. I checked in breeze with datafusion 51.0.0: `GoogleCloud(bucket_name="b", service_account_path="/nonexistent/key.json")` sails past `except Exception` here and in `_register_object_store`, so the `ObjectStoreCreationException` wrapping never fires for the most likely misconfig (the `finally` still deletes the temp file, so that part is fine). Checking `Path(key_path).is_file()` before this call and raising `ObjectStoreCreationException` yourself would give users a normal error instead of a Rust panic. Related: `test_gcs_provider_failure` models a plain `Exception`, which the real binding does not raise.



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -169,6 +169,26 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, Any]:
                 credentials = self._remove_none_values(credentials)
                 extra_config = _fetch_extra_configs(["region", "endpoint"])
 
+            case "google_cloud_platform":
+                try:
+                    from airflow.providers.google.common.hooks.base_google import get_field
+                except ImportError:
+                    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+                    raise AirflowOptionalProviderFeatureException(
+                        "Failed to import get_field. To use the GCS storage functionality, please install "
+                        "the apache-airflow-providers-google package."
+                    )
+                extra_dejson = conn.extra_dejson
+                key_path = get_field(extra_dejson, "key_path")
+                keyfile_dict = get_field(extra_dejson, "keyfile_dict")
+                if key_path and keyfile_dict:
+                    raise ValueError(
+                        "The `keyfile_dict` and `key_path` fields are mutually exclusive. "
+                        "Please provide only one value."
+                    )
+                credentials = self._remove_none_values({"key_path": key_path, "keyfile_dict": keyfile_dict})

Review Comment:
   One correction to the ADC framing. With `service_account_path=None` the binding builds a bare `GoogleCloudStorageBuilder` and never calls `from_env()` (the S3 binding does call `AmazonS3Builder::from_env()`), so `GOOGLE_APPLICATION_CREDENTIALS` is not read at all. What the fallback actually picks up is `$HOME/.config/gcloud/application_default_credentials.json` or the GCE/GKE metadata server, and a worker that relies on the env var will build fine here and only fail at query time. Either honouring `GOOGLE_APPLICATION_CREDENTIALS` as the `key_path` fallback when it points at a `service_account` file, or documenting the narrower behaviour, would avoid that surprise.



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -169,6 +169,26 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, Any]:
                 credentials = self._remove_none_values(credentials)
                 extra_config = _fetch_extra_configs(["region", "endpoint"])
 
+            case "google_cloud_platform":
+                try:
+                    from airflow.providers.google.common.hooks.base_google import get_field
+                except ImportError:
+                    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+                    raise AirflowOptionalProviderFeatureException(
+                        "Failed to import get_field. To use the GCS storage functionality, please install "
+                        "the apache-airflow-providers-google package."
+                    )
+                extra_dejson = conn.extra_dejson
+                key_path = get_field(extra_dejson, "key_path")
+                keyfile_dict = get_field(extra_dejson, "keyfile_dict")

Review Comment:
   A `google_cloud_platform` connection configured with `key_secret_name`, `credential_config_file`, or `impersonation_chain` works with every Google operator, but here those fields are silently dropped and the store falls through to the metadata-server identity. Raising a `ValueError` naming the unsupported field seems safer than quietly querying GCS as a different principal than the connection says.



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -53,6 +57,53 @@ def get_scheme(self) -> str:
         return "s3://"
 
 
+class GCSObjectStorageProvider(ObjectStorageProvider):
+    """GCS Object Storage Provider using DataFusion's GoogleCloud."""

Review Comment:
   `providers/common/sql/docs/operators.rst` still lists GCS under "not yet supported" in the Supported Storage Systems section, and nothing yet tells users which connection fields this reads (`key_path`, `keyfile_dict`, or what happens when neither is set). A short GCS Storage section next to the S3 one would close both gaps.



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -169,6 +169,26 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, Any]:
                 credentials = self._remove_none_values(credentials)
                 extra_config = _fetch_extra_configs(["region", "endpoint"])
 
+            case "google_cloud_platform":
+                try:
+                    from airflow.providers.google.common.hooks.base_google import get_field

Review Comment:
   Is the google provider dependency worth it for `get_field`? It is a ten-line prefix lookup for the `extra__google_cloud_platform__` back-compat names, and unlike the `aws` branch, where `AwsGenericHook.get_credentials()` does the actual credential resolution, nothing here needs the google stack because DataFusion does all the GCS I/O itself. Inlining that lookup would let `common-sql[datafusion]` users query GCS without installing `google-cloud-*`. If the plan is to lean on `GoogleBaseHook` properly (resolving `key_secret_name`, for example) then the dependency makes sense, but the branch does not do that today.



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Add GCS support to the DataFusion object storage layer [airflow]

Posted by "pankajastro (via GitHub)" <gi...@apache.org>.
pankajastro commented on code in PR #73370:
URL: https://github.com/apache/airflow/pull/73370#discussion_r4056501805


##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -53,6 +57,53 @@ def get_scheme(self) -> str:
         return "s3://"
 
 
+class GCSObjectStorageProvider(ObjectStorageProvider):
+    """GCS Object Storage Provider using DataFusion's GoogleCloud."""
+
+    @property
+    def get_storage_type(self) -> StorageType:
+        """Return the storage type."""
+        return StorageType.GCS
+
+    def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None):
+        """Create a GCS object store using DataFusion's GoogleCloud."""
+        if connection_config is None:
+            raise ValueError(f"connection_config must be provided for {self.get_storage_type}")
+
+        credentials = connection_config.credentials
+        key_path = credentials.get("key_path")
+        keyfile_dict = credentials.get("keyfile_dict")
+        temp_key_path: str | None = None
+
+        try:
+            bucket = self.get_bucket(path)
+
+            if not key_path and keyfile_dict:
+                # DataFusion's GoogleCloud binding only accepts a file path, not inline JSON,
+                # so materialize keyfile_dict to a temp file. The credentials file is read once
+                # at construction time and never touched again, so it's safe to delete right after.
+                key_content = keyfile_dict if isinstance(keyfile_dict, str) else json.dumps(keyfile_dict)
+                with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as key_file:
+                    key_file.write(key_content)
+                temp_key_path = key_path = key_file.name
+
+            gcs_store = GoogleCloud(bucket_name=bucket, service_account_path=key_path)

Review Comment:
   Fixed — pre-flight `Path.is_file()` check plus widened to `except BaseException` so the panic is caught cleanly. Updated `test_gcs_provider_failure` and added a real (unmocked) repro test per your note.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -169,6 +169,26 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, Any]:
                 credentials = self._remove_none_values(credentials)
                 extra_config = _fetch_extra_configs(["region", "endpoint"])
 
+            case "google_cloud_platform":
+                try:
+                    from airflow.providers.google.common.hooks.base_google import get_field
+                except ImportError:
+                    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+                    raise AirflowOptionalProviderFeatureException(
+                        "Failed to import get_field. To use the GCS storage functionality, please install "
+                        "the apache-airflow-providers-google package."
+                    )
+                extra_dejson = conn.extra_dejson
+                key_path = get_field(extra_dejson, "key_path")
+                keyfile_dict = get_field(extra_dejson, "keyfile_dict")
+                if key_path and keyfile_dict:
+                    raise ValueError(
+                        "The `keyfile_dict` and `key_path` fields are mutually exclusive. "
+                        "Please provide only one value."
+                    )
+                credentials = self._remove_none_values({"key_path": key_path, "keyfile_dict": keyfile_dict})

Review Comment:
   Good catch — added an explicit `GOOGLE_APPLICATION_CREDENTIALS` fallback (only when neither `key_path` nor `keyfile_dict` is set), instead of relying on the binding's narrower built-in behavior.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -169,6 +169,26 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, Any]:
                 credentials = self._remove_none_values(credentials)
                 extra_config = _fetch_extra_configs(["region", "endpoint"])
 
+            case "google_cloud_platform":
+                try:
+                    from airflow.providers.google.common.hooks.base_google import get_field
+                except ImportError:
+                    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+                    raise AirflowOptionalProviderFeatureException(
+                        "Failed to import get_field. To use the GCS storage functionality, please install "
+                        "the apache-airflow-providers-google package."
+                    )
+                extra_dejson = conn.extra_dejson
+                key_path = get_field(extra_dejson, "key_path")
+                keyfile_dict = get_field(extra_dejson, "keyfile_dict")

Review Comment:
   Added — raises `ValueError` naming the field for `key_secret_name`, `credential_config_file`, or `impersonation_chain` instead of falling through to metadata-server identity.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -169,6 +169,26 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str, Any]:
                 credentials = self._remove_none_values(credentials)
                 extra_config = _fetch_extra_configs(["region", "endpoint"])
 
+            case "google_cloud_platform":
+                try:
+                    from airflow.providers.google.common.hooks.base_google import get_field

Review Comment:
   Agreed — dropped the dependency entirely and inlined the trivial extras lookup. Nothing here needs the SDK stack.
   
   ---
   Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting



##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py:
##########
@@ -53,6 +57,53 @@ def get_scheme(self) -> str:
         return "s3://"
 
 
+class GCSObjectStorageProvider(ObjectStorageProvider):
+    """GCS Object Storage Provider using DataFusion's GoogleCloud."""

Review Comment:
   Added a GCS Storage section to `operators.rst` with the credential priority order and an example.
   
   ---
   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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Add GCS support to the DataFusion object storage layer [airflow]

Posted by "pankajastro (via GitHub)" <gi...@apache.org>.
pankajastro merged PR #73370:
URL: https://github.com/apache/airflow/pull/73370


-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Add GCS support to the DataFusion object storage layer [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #73370:
URL: https://github.com/apache/airflow/pull/73370#issuecomment-5749487786

   Two notes on the description:
   
   - "falling back to Application Default Credentials when neither is set": with `service_account_path=None` the binding builds a bare `GoogleCloudStorageBuilder` and never calls `from_env()`, so `GOOGLE_APPLICATION_CREDENTIALS` is not read. The fallback only finds the gcloud ADC file under `$HOME/.config/gcloud/` or the GCE/GKE metadata server. Either tighten the wording or honour the env var as the `key_path` fallback.
   - "docs updated to mention GCS alongside S3": `providers/common/sql/docs/operators.rst` still lists GCS under "not yet supported" in the Supported Storage Systems section (line 302 at this head).
   


-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org