potiuk commented on code in PR #64941: URL: https://github.com/apache/airflow/pull/64941#discussion_r3979771527
########## providers/amazon/src/airflow/providers/amazon/aws/datafusion/object_storage.py: ########## @@ -0,0 +1,66 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from datafusion.object_store import AmazonS3 + +from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook +from airflow.providers.common.compat.sdk import BaseHook +from airflow.providers.common.sql.config import ConnectionConfig, StorageType +from airflow.providers.common.sql.datafusion.base import ObjectStorageProvider +from airflow.providers.common.sql.datafusion.exceptions import ObjectStoreCreationException + + +class S3ObjectStorageProvider(ObjectStorageProvider): + """S3 Object Storage Provider using DataFusion's AmazonS3.""" + + @property + def get_storage_type(self) -> StorageType: + """Return the storage type.""" + return StorageType.S3 + + def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None): + """Create an S3 object store using DataFusion's AmazonS3.""" + if connection_config is None: + raise ValueError(f"connection_config must be provided for {self.get_storage_type}") + + try: + conn = BaseHook.get_connection(connection_config.conn_id) + aws_hook: AwsGenericHook = AwsGenericHook(aws_conn_id=conn.conn_id, client_type="s3") + creds = aws_hook.get_credentials() + + credentials = { + "access_key_id": conn.login or creds.access_key, + "secret_access_key": conn.password or creds.secret_key, + "session_token": creds.token if creds.token else None, + } + credentials = {k: v for k, v in credentials.items() if v is not None} + extra_config = {k: conn.extra_dejson[k] for k in ["region", "endpoint"] if k in conn.extra_dejson} Review Comment: This block changes the contract of `create_object_store`, which a relocation should not do. On `main` the implementation used what the caller supplied: ```python credentials = connection_config.credentials bucket = self.get_bucket(path) s3_store = AmazonS3(**credentials, **connection_config.extra_config, bucket_name=bucket) ``` Here, `connection_config.credentials` and `connection_config.extra_config` are never read. The method resolves an Airflow connection from `conn_id`, builds an `AwsGenericHook`, and derives credentials from the connection plus `extra_dejson["region"]` / `["endpoint"]`. Both fields still exist on the public `ConnectionConfig` dataclass in `common/sql/config.py` — for S3 they are now dead. Two ways this bites an existing caller who passes credentials explicitly: 1. No stored connection matching `conn_id` — `BaseHook.get_connection` raises, gets wrapped by the `except Exception` below, and surfaces as `ObjectStoreCreationException`. Previously this worked without any Airflow connection at all. 2. A stored connection does exist — the object store is created with **different** credentials and endpoint than the caller asked for. That failure is silent, which is the worse of the two. The deprecation shim compounds it. `object_storage_provider.__getattr__` warns that the import moved and then returns *this* class, so code following the deprecation notice gets new behaviour under the banner of a compatibility path. A shim that changes semantics is not a shim. What makes me confident this was not a deliberate, considered change: `test_s3_provider_success` in `common/sql/tests/.../test_object_storage_provider.py` pinned the old contract exactly — ```python credentials={"access_key_id": "fake_key", "secret_access_key": "fake_secret"}, ... AmazonS3(access_key_id="fake_key", secret_access_key="fake_secret", bucket_name="demo-data") ``` — and it was deleted rather than adapted. The replacement tests in the amazon package mock `BaseHook.get_connection` and `AwsGenericHook` and assert the new shape, so CI is green and the contract change leaves no trace in the diff's test signal. Deriving credentials through `AwsGenericHook` may well be the better long-term design — it picks up the standard AWS credential chain. But it is a separate, user-visible behaviour change, and it should not ride along inside a decoupling PR. Either: - preserve the existing contract here (honour `connection_config.credentials` / `extra_config` when supplied, fall back to the hook only when they are empty), restore a regression test that exercises the explicit-credentials path through the deprecated import, and leave the redesign to its own PR; or - make the change deliberately, in which case it needs a changelog entry in `providers/amazon/docs/changelog.rst`, the same for common-sql, and the `ConnectionConfig` fields either removed or documented as ignored for S3. Either way the current state — changed behaviour, deleted test, unchanged docstring, "compatible" shim — should not merge. This finding came from the adversarial second read on this PR, not my own pass; I verified it against `git show main:` and the deleted test before including it. ########## providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py: ########## @@ -70,18 +44,58 @@ def get_scheme(self) -> str: return "file://" +_STORAGE_TYPE_PROVIDER_HINTS: dict[str, str] = { + "s3": "apache-airflow-providers-amazon[datafusion]", +} + + +def _missing_provider_message(type_key: str) -> str: + hint = _STORAGE_TYPE_PROVIDER_HINTS.get(type_key, "the appropriate provider package") + return f"No ObjectStorageProvider registered for storage type '{type_key}'. Install or upgrade {hint}." + + +def _get_legacy_object_storage_provider(type_key: str) -> ObjectStorageProvider: + if type_key == StorageType.S3.value: + try: + from airflow.providers.amazon.aws.datafusion.object_storage import S3ObjectStorageProvider + except ImportError as err: + raise ValueError(_missing_provider_message(type_key)) from err + return S3ObjectStorageProvider() + + raise ValueError(_missing_provider_message(type_key)) + + def get_object_storage_provider(storage_type: StorageType) -> ObjectStorageProvider: """Get an object storage provider based on the storage type.""" - # TODO: Add support for GCS, Azure, HTTP: https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html - providers: dict[StorageType, type] = { - StorageType.S3: S3ObjectStorageProvider, - StorageType.LOCAL: LocalObjectStorageProvider, - } - - if storage_type not in providers: - raise ValueError( - f"Unsupported storage type: {storage_type}. Supported types: {list(providers.keys())}" + if storage_type == StorageType.LOCAL: + return LocalObjectStorageProvider() + + type_key = storage_type.value + + from airflow.providers_manager import ProvidersManager + + manager = ProvidersManager() + if not hasattr(manager, "object_storage_providers"): + return _get_legacy_object_storage_provider(type_key) + + registry = manager.object_storage_providers + if type_key in registry: + provider_cls = import_string(registry[type_key].provider_class_name) + return provider_cls() Review Comment: `import_string` is unguarded here, so the helpful message this module goes to the trouble of building never reaches the most likely failure. Compare with the legacy branch a few lines up, which does guard it: ```python try: from airflow.providers.amazon.aws.datafusion.object_storage import S3ObjectStorageProvider except ImportError as err: raise ValueError(_missing_provider_message(type_key)) from err ``` `_missing_provider_message` resolves to *"No ObjectStorageProvider registered for storage type 's3'. Install or upgrade apache-airflow-providers-amazon[datafusion]."* — exactly the guidance a user needs. On the registry path they get a bare `ModuleNotFoundError: No module named 'datafusion'` instead. The gap is reachable in an ordinary setup, because registration and importability are decoupled: - `providers/amazon/provider.yaml` declares the `object-storage-providers` entry unconditionally, so `s3` lands in `ProvidersManager.object_storage_providers` whenever the amazon provider is installed at all. - `airflow/providers/amazon/aws/datafusion/object_storage.py:19` imports `datafusion.object_store` at module scope, and `datafusion` is an **optional** amazon extra. So "amazon installed without `[datafusion]`" — which nothing in the packaging prevents, see the thread on `common/sql/pyproject.toml` — puts `s3` in the registry with a class that cannot be imported, and this line turns that into a traceback rather than the install hint. The result is inverted: the old-Airflow compatibility path gives the good error, the modern path gives the raw one. ```suggestion registry = manager.object_storage_providers if type_key in registry: try: provider_cls = import_string(registry[type_key].provider_class_name) except ImportError as err: raise ValueError(_missing_provider_message(type_key)) from err return provider_cls() ``` Worth a test for it too — registry entry present, target module unimportable, assert the `ValueError` and its hint rather than an `ImportError`. That is the case users will actually hit. For context: this is the same concern raised earlier in review about `import_string` failing for a registered storage type. That thread is marked resolved, but only the legacy branch was guarded; this one still is not. ########## providers/common/sql/pyproject.toml: ########## @@ -100,6 +97,9 @@ dependencies = [ "apache.iceberg" = [ "apache-airflow-providers-apache-iceberg" ] +"amazon" = [ + "apache-airflow-providers-amazon" +] Review Comment: This extra needs a version floor, and right now it is the one line in this file that changed without changing anything. The diff deletes `"amazon" = [...]` from its old position and re-adds it here — a pure reorder, net diff zero. So the file was touched, but the dependency is still unversioned `apache-airflow-providers-amazon`, which any already-installed older release satisfies. That matters because this PR relocates the implementation. Upgrade `common-sql` while leaving `amazon` at an older release and S3 object storage stops working, on both branches of `get_object_storage_provider`: - **Older Airflow core** takes the `_get_legacy_object_storage_provider` path and tries to import `airflow.providers.amazon.aws.datafusion.object_storage`, which does not exist in the older amazon package. - **Newer Airflow core** consults `ProvidersManager.object_storage_providers`, and the older amazon's `provider.yaml` declares no `object-storage-providers` entry, so there is no `s3` key to find. Both end in a `ValueError` for a configuration that worked before the upgrade. The error message is at least the friendly one here, but a previously-working deployment breaking on a single-provider upgrade is the kind of thing the dependency floor exists to prevent. The fix is to pin the floor to the amazon release that first contains the relocated module, using the project's usual convention for referencing the next release. Please also settle the related question raised earlier in review: since the S3 provider class imports `datafusion.object_store` at module scope, whether this should be `apache-airflow-providers-amazon[datafusion]` rather than plain `apache-airflow-providers-amazon` — see the thread on `object_storage_provider.py:83`, where the missing extra produces a bare `ModuleNotFoundError`. That earlier thread was closed by updating the PR description; the packaging question underneath it is still open. If the reorder was not intentional, dropping it would leave this file out of the diff entirely — which would itself be a signal that the compatibility question has not been handled. This finding came from the adversarial second read; the observation that the change is a no-op reorder is from my own pass. -- 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]
