From 1ab983a808b3c2255e8c5c92be08e35e30b17f7f Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Thu, 17 Oct 2024 17:12:50 +0800 Subject: [PATCH] feat(providers/amazon): Use asset in common provider --- generated/provider_dependencies.json | 2 +- .../airflow/providers/amazon/aws/assets/s3.py | 14 +-- .../airflow/providers/amazon/aws/hooks/s3.py | 106 +++++++++++++----- .../utils/asset_compat_lineage_collector.py | 106 ------------------ .../airflow/providers/amazon/provider.yaml | 2 +- .../aws/auth_manager/test_aws_auth_manager.py | 8 +- providers/tests/amazon/aws/hooks/test_s3.py | 2 +- 7 files changed, 89 insertions(+), 151 deletions(-) delete mode 100644 providers/src/airflow/providers/amazon/aws/utils/asset_compat_lineage_collector.py diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index fb50939ac7fb8..2cd9966b7c757 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -26,7 +26,7 @@ "amazon": { "deps": [ "PyAthena>=3.0.10", - "apache-airflow-providers-common-compat>=1.1.0", + "apache-airflow-providers-common-compat>=1.2.1", "apache-airflow-providers-common-sql>=1.3.1", "apache-airflow-providers-http", "apache-airflow>=2.8.0", diff --git a/providers/src/airflow/providers/amazon/aws/assets/s3.py b/providers/src/airflow/providers/amazon/aws/assets/s3.py index 378e7e977ab56..4d02b156afb0c 100644 --- a/providers/src/airflow/providers/amazon/aws/assets/s3.py +++ b/providers/src/airflow/providers/amazon/aws/assets/s3.py @@ -19,16 +19,14 @@ from typing import TYPE_CHECKING from airflow.providers.amazon.aws.hooks.s3 import S3Hook - -try: - from airflow.assets import Asset -except ModuleNotFoundError: - from airflow.datasets import Dataset as Asset # type: ignore[no-redef] +from airflow.providers.common.compat.assets import Asset if TYPE_CHECKING: from urllib.parse import SplitResult - from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset + from airflow.providers.common.compat.openlineage.facet import ( + Dataset as OpenLineageDataset, + ) def create_asset(*, bucket: str, key: str, extra=None) -> Asset: @@ -43,7 +41,9 @@ def sanitize_uri(uri: SplitResult) -> SplitResult: def convert_asset_to_openlineage(asset: Asset, lineage_context) -> OpenLineageDataset: """Translate Asset with valid AIP-60 uri to OpenLineage with assistance from the hook.""" - from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset + from airflow.providers.common.compat.openlineage.facet import ( + Dataset as OpenLineageDataset, + ) bucket, key = S3Hook.parse_s3_url(asset.uri) return OpenLineageDataset(namespace=f"s3://{bucket}", name=key if key else "/") diff --git a/providers/src/airflow/providers/amazon/aws/hooks/s3.py b/providers/src/airflow/providers/amazon/aws/hooks/s3.py index 6efb5953b8eee..0ede24bc6dd5d 100644 --- a/providers/src/airflow/providers/amazon/aws/hooks/s3.py +++ b/providers/src/airflow/providers/amazon/aws/hooks/s3.py @@ -41,14 +41,16 @@ from uuid import uuid4 if TYPE_CHECKING: - from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject + from mypy_boto3_s3.service_resource import ( + Bucket as S3Bucket, + Object as S3ResourceObject, + ) from airflow.utils.types import ArgNotSet with suppress(ImportError): from aiobotocore.client import AioBaseClient -from importlib.util import find_spec from asgiref.sync import sync_to_async from boto3.s3.transfer import S3Transfer, TransferConfig @@ -58,15 +60,9 @@ from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.utils.tags import format_tags +from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector from airflow.utils.helpers import chunks -if find_spec("airflow.assets"): - from airflow.lineage.hook import get_hook_lineage_collector -else: - # TODO: import from common.compat directly after common.compat providers with - # asset_compat_lineage_collector released - from airflow.providers.amazon.aws.utils.asset_compat_lineage_collector import get_hook_lineage_collector - logger = logging.getLogger(__name__) @@ -148,9 +144,10 @@ def wrapper(*args, **kwargs) -> Callable: if "bucket_name" not in bound_args.arguments: with suppress(S3HookUriParseFailure): - bound_args.arguments["bucket_name"], bound_args.arguments[key_name] = S3Hook.parse_s3_url( - bound_args.arguments[key_name] - ) + ( + bound_args.arguments["bucket_name"], + bound_args.arguments[key_name], + ) = S3Hook.parse_s3_url(bound_args.arguments[key_name]) return func(*bound_args.args, **bound_args.kwargs) @@ -318,7 +315,8 @@ def check_for_bucket(self, bucket_name: str | None = None) -> bool: self.log.info('Bucket "%s" does not exist', bucket_name) elif return_code == 403: self.log.error( - 'Access to bucket "%s" is forbidden or there was an error with the request', bucket_name + 'Access to bucket "%s" is forbidden or there was an error with the request', + bucket_name, ) self.log.error(e) return False @@ -359,7 +357,8 @@ def create_bucket(self, bucket_name: str | None = None, region_name: str | None self.get_conn().create_bucket(Bucket=bucket_name) else: self.get_conn().create_bucket( - Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region_name} + Bucket=bucket_name, + CreateBucketConfiguration={"LocationConstraint": region_name}, ) @provide_bucket_name @@ -410,7 +409,10 @@ def list_prefixes( paginator = self.get_conn().get_paginator("list_objects_v2") response = paginator.paginate( - Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config + Bucket=bucket_name, + Prefix=prefix, + Delimiter=delimiter, + PaginationConfig=config, ) prefixes: list[str] = [] @@ -471,7 +473,10 @@ async def list_prefixes_async( paginator = client.get_paginator("list_objects_v2") response = paginator.paginate( - Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config + Bucket=bucket_name, + Prefix=prefix, + Delimiter=delimiter, + PaginationConfig=config, ) prefixes = [] @@ -569,7 +574,11 @@ async def check_key_async( return await self._check_key_async(client, bucket, wildcard_match, bucket_keys, use_regex) async def check_for_prefix_async( - self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None + self, + client: AioBaseClient, + prefix: str, + delimiter: str, + bucket_name: str | None = None, ) -> bool: """ Check that a prefix exists in a bucket. @@ -587,7 +596,11 @@ async def check_for_prefix_async( return prefix in plist async def _check_for_prefix_async( - self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None + self, + client: AioBaseClient, + prefix: str, + delimiter: str, + bucket_name: str | None = None, ) -> bool: return await self.check_for_prefix_async( client, prefix=prefix, delimiter=delimiter, bucket_name=bucket_name @@ -643,7 +656,10 @@ async def _list_keys_async( paginator = client.get_paginator("list_objects_v2") response = paginator.paginate( - Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config + Bucket=bucket_name, + Prefix=prefix, + Delimiter=delimiter, + PaginationConfig=config, ) keys = [] @@ -655,7 +671,10 @@ async def _list_keys_async( return keys def _list_key_object_filter( - self, keys: list, from_datetime: datetime | None = None, to_datetime: datetime | None = None + self, + keys: list, + from_datetime: datetime | None = None, + to_datetime: datetime | None = None, ) -> list: def _is_in_period(input_date: datetime) -> bool: if from_datetime is not None and input_date <= from_datetime: @@ -766,7 +785,10 @@ async def is_keys_unchanged_async( "message": success_message, } - self.log.error("FAILURE: Inactivity Period passed, not enough objects found in %s", path) + self.log.error( + "FAILURE: Inactivity Period passed, not enough objects found in %s", + path, + ) return { "status": "error", "message": f"FAILURE: Inactivity Period passed, not enough objects found in {path}", @@ -1109,7 +1131,13 @@ def load_file( extra_args["ACL"] = acl_policy client = self.get_conn() - client.upload_file(filename, bucket_name, key, ExtraArgs=extra_args, Config=self.transfer_config) + client.upload_file( + filename, + bucket_name, + key, + ExtraArgs=extra_args, + Config=self.transfer_config, + ) get_hook_lineage_collector().add_input_asset( context=self, scheme="file", asset_kwargs={"path": filename} ) @@ -1308,18 +1336,32 @@ def copy_object( ) source_bucket_name, source_bucket_key = self.get_s3_bucket_key( - source_bucket_name, source_bucket_key, "source_bucket_name", "source_bucket_key" + source_bucket_name, + source_bucket_key, + "source_bucket_name", + "source_bucket_key", ) - copy_source = {"Bucket": source_bucket_name, "Key": source_bucket_key, "VersionId": source_version_id} + copy_source = { + "Bucket": source_bucket_name, + "Key": source_bucket_key, + "VersionId": source_version_id, + } response = self.get_conn().copy_object( - Bucket=dest_bucket_name, Key=dest_bucket_key, CopySource=copy_source, **kwargs + Bucket=dest_bucket_name, + Key=dest_bucket_key, + CopySource=copy_source, + **kwargs, ) get_hook_lineage_collector().add_input_asset( - context=self, scheme="s3", asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key} + context=self, + scheme="s3", + asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}, ) get_hook_lineage_collector().add_output_asset( - context=self, scheme="s3", asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key} + context=self, + scheme="s3", + asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}, ) return response @@ -1435,7 +1477,10 @@ def download_file( file_path = Path(local_dir, subdir, filename_in_s3) if file_path.is_file(): - self.log.error("file '%s' already exists. Failing the task and not overwriting it", file_path) + self.log.error( + "file '%s' already exists. Failing the task and not overwriting it", + file_path, + ) raise FileExistsError file_path.parent.mkdir(exist_ok=True, parents=True) @@ -1484,7 +1529,10 @@ def generate_presigned_url( s3_client = self.get_conn() try: return s3_client.generate_presigned_url( - ClientMethod=client_method, Params=params, ExpiresIn=expires_in, HttpMethod=http_method + ClientMethod=client_method, + Params=params, + ExpiresIn=expires_in, + HttpMethod=http_method, ) except ClientError as e: diff --git a/providers/src/airflow/providers/amazon/aws/utils/asset_compat_lineage_collector.py b/providers/src/airflow/providers/amazon/aws/utils/asset_compat_lineage_collector.py deleted file mode 100644 index 50fbc3d0996aa..0000000000000 --- a/providers/src/airflow/providers/amazon/aws/utils/asset_compat_lineage_collector.py +++ /dev/null @@ -1,106 +0,0 @@ -# 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 importlib.util import find_spec - - -def _get_asset_compat_hook_lineage_collector(): - from airflow.lineage.hook import get_hook_lineage_collector - - collector = get_hook_lineage_collector() - - if all( - getattr(collector, asset_method_name, None) - for asset_method_name in ("add_input_asset", "add_output_asset", "collected_assets") - ): - return collector - - # dataset is renamed as asset in Airflow 3.0 - - from functools import wraps - - from airflow.lineage.hook import DatasetLineageInfo, HookLineage - - DatasetLineageInfo.asset = DatasetLineageInfo.dataset - - def rename_dataset_kwargs_as_assets_kwargs(function): - @wraps(function) - def wrapper(*args, **kwargs): - if "asset_kwargs" in kwargs: - kwargs["dataset_kwargs"] = kwargs.pop("asset_kwargs") - - if "asset_extra" in kwargs: - kwargs["dataset_extra"] = kwargs.pop("asset_extra") - - return function(*args, **kwargs) - - return wrapper - - collector.create_asset = rename_dataset_kwargs_as_assets_kwargs(collector.create_dataset) - collector.add_input_asset = rename_dataset_kwargs_as_assets_kwargs(collector.add_input_dataset) - collector.add_output_asset = rename_dataset_kwargs_as_assets_kwargs(collector.add_output_dataset) - - def collected_assets_compat(collector) -> HookLineage: - """Get the collected hook lineage information.""" - lineage = collector.collected_datasets - return HookLineage( - [ - DatasetLineageInfo(dataset=item.dataset, count=item.count, context=item.context) - for item in lineage.inputs - ], - [ - DatasetLineageInfo(dataset=item.dataset, count=item.count, context=item.context) - for item in lineage.outputs - ], - ) - - setattr( - collector.__class__, - "collected_assets", - property(lambda collector: collected_assets_compat(collector)), - ) - - return collector - - -def get_hook_lineage_collector(): - # HookLineageCollector added in 2.10 - try: - if find_spec("airflow.assets"): - # Dataset has been renamed as Asset in 3.0 - from airflow.lineage.hook import get_hook_lineage_collector - - return get_hook_lineage_collector() - - return _get_asset_compat_hook_lineage_collector() - except ImportError: - - class NoOpCollector: - """ - NoOpCollector is a hook lineage collector that does nothing. - - It is used when you want to disable lineage collection. - """ - - def add_input_asset(self, *_, **__): - pass - - def add_output_asset(self, *_, **__): - pass - - return NoOpCollector() diff --git a/providers/src/airflow/providers/amazon/provider.yaml b/providers/src/airflow/providers/amazon/provider.yaml index 71d097f298198..a8d182b904954 100644 --- a/providers/src/airflow/providers/amazon/provider.yaml +++ b/providers/src/airflow/providers/amazon/provider.yaml @@ -93,7 +93,7 @@ versions: dependencies: - apache-airflow>=2.8.0 - - apache-airflow-providers-common-compat>=1.1.0 + - apache-airflow-providers-common-compat>=1.2.1 - apache-airflow-providers-common-sql>=1.3.1 - apache-airflow-providers-http # We should update minimum version of boto3 and here regularly to avoid `pip` backtracking with the number diff --git a/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py b/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py index c8c1bb03d4bb6..b62fa02889bce 100644 --- a/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py +++ b/providers/tests/amazon/aws/auth_manager/test_aws_auth_manager.py @@ -68,12 +68,8 @@ from airflow.auth.managers.models.resource_details import AssetDetails from airflow.security.permissions import RESOURCE_ASSET else: - try: - from airflow.auth.managers.models.resource_details import AssetDetails - from airflow.security.permissions import RESOURCE_ASSET - except ImportError: - from airflow.auth.managers.models.resource_details import DatasetDetails as AssetDetails - from airflow.security.permissions import RESOURCE_DATASET as RESOURCE_ASSET + from airflow.providers.common.compat.assets import AssetDetails + from airflow.providers.common.compat.security.permissions import RESOURCE_ASSET pytestmark = [ pytest.mark.skipif(not AIRFLOW_V_2_9_PLUS, reason="Test requires Airflow 2.9+"), diff --git a/providers/tests/amazon/aws/hooks/test_s3.py b/providers/tests/amazon/aws/hooks/test_s3.py index 727322299f174..3735c3dc0385d 100644 --- a/providers/tests/amazon/aws/hooks/test_s3.py +++ b/providers/tests/amazon/aws/hooks/test_s3.py @@ -64,7 +64,7 @@ def s3_bucket(mocked_s3_res): @pytest.fixture def hook_lineage_collector(): from airflow.lineage import hook - from airflow.providers.amazon.aws.hooks.s3 import get_hook_lineage_collector + from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector hook._hook_lineage_collector = None hook._hook_lineage_collector = hook.HookLineageCollector()