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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
