cruseakshay commented on code in PR #64941:
URL: https://github.com/apache/airflow/pull/64941#discussion_r4016152955
##########
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:
applied your suggestion. A registered provider whose module can't be
imported now raises the hinted ValueError.
--
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]
