You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by mo...@apache.org on 2024/07/15 23:48:33 UTC
(airflow) branch aip-62/s3hook updated (ba1bb1268a -> a8e0189a2f)
This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a change to branch aip-62/s3hook
in repository https://gitbox.apache.org/repos/asf/airflow.git
discard ba1bb1268a openlineage: add support for hook lineage for S3Hook
new a8e0189a2f openlineage: add support for hook lineage for S3Hook
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (ba1bb1268a)
\
N -- N -- N refs/heads/aip-62/s3hook (a8e0189a2f)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../amazon/aws/datasets}/__init__.py | 0
.../amazon/aws/{utils/rds.py => datasets/s3.py} | 9 +++------
.../providers/amazon/aws/datasets}/__init__.py | 0
.../amazon/aws/datasets/test_s3.py} | 17 ++++++-----------
4 files changed, 9 insertions(+), 17 deletions(-)
copy airflow/{api_connexion => providers/amazon/aws/datasets}/__init__.py (100%)
copy airflow/providers/amazon/aws/{utils/rds.py => datasets/s3.py} (85%)
copy {airflow/api_connexion => tests/providers/amazon/aws/datasets}/__init__.py (100%)
copy tests/{core/test_example_dags_system.py => providers/amazon/aws/datasets/test_s3.py} (68%)
(airflow) 01/01: openlineage: add support for hook lineage for S3Hook
Posted by mo...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a commit to branch aip-62/s3hook
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit a8e0189a2ffd18cc33e7c0255ea82f13aadfef05
Author: Maciej Obuchowski <ob...@gmail.com>
AuthorDate: Tue Jul 16 01:27:07 2024 +0200
openlineage: add support for hook lineage for S3Hook
Signed-off-by: Maciej Obuchowski <ob...@gmail.com>
---
airflow/providers/amazon/aws/datasets/__init__.py | 16 +++++++
airflow/providers/amazon/aws/datasets/s3.py | 23 +++++++++
airflow/providers/amazon/aws/hooks/s3.py | 24 +++++++++-
airflow/providers/amazon/provider.yaml | 1 +
airflow/providers_manager.py | 5 +-
generated/provider_dependencies.json | 1 +
tests/conftest.py | 10 ++++
tests/providers/amazon/aws/datasets/__init__.py | 16 +++++++
tests/providers/amazon/aws/datasets/test_s3.py | 27 +++++++++++
tests/providers/amazon/aws/hooks/test_s3.py | 58 +++++++++++++++++++++--
10 files changed, 175 insertions(+), 6 deletions(-)
diff --git a/airflow/providers/amazon/aws/datasets/__init__.py b/airflow/providers/amazon/aws/datasets/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/airflow/providers/amazon/aws/datasets/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/airflow/providers/amazon/aws/datasets/s3.py b/airflow/providers/amazon/aws/datasets/s3.py
new file mode 100644
index 0000000000..b89cdefe36
--- /dev/null
+++ b/airflow/providers/amazon/aws/datasets/s3.py
@@ -0,0 +1,23 @@
+# 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 airflow.datasets import Dataset
+
+
+def create_dataset(*, bucket: str, key: str) -> Dataset:
+ return Dataset(uri=f"s3://{bucket}/{key}")
diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
index 8ca93766e2..d0f6e027f4 100644
--- a/airflow/providers/amazon/aws/hooks/s3.py
+++ b/airflow/providers/amazon/aws/hooks/s3.py
@@ -41,6 +41,8 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Callable
from urllib.parse import urlsplit
from uuid import uuid4
+from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
+
if TYPE_CHECKING:
from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject
@@ -1111,6 +1113,10 @@ class S3Hook(AwsBaseHook):
client = self.get_conn()
client.upload_file(filename, bucket_name, key, ExtraArgs=extra_args, Config=self.transfer_config)
+ get_hook_lineage_collector().add_input_dataset({"scheme": "file", "path": filename}, self)
+ get_hook_lineage_collector().add_output_dataset(
+ {"scheme": "s3", "bucket": bucket_name, "key": key}, self
+ )
@unify_bucket_name_and_key
@provide_bucket_name
@@ -1251,6 +1257,10 @@ class S3Hook(AwsBaseHook):
ExtraArgs=extra_args,
Config=self.transfer_config,
)
+ # No input because file_obj can be anything - handle in calling function if possible
+ get_hook_lineage_collector().add_output_dataset(
+ {"scheme": "s3", "bucket": bucket_name, "key": key}, self
+ )
def copy_object(
self,
@@ -1306,6 +1316,12 @@ class S3Hook(AwsBaseHook):
response = self.get_conn().copy_object(
Bucket=dest_bucket_name, Key=dest_bucket_key, CopySource=copy_source, **kwargs
)
+ get_hook_lineage_collector().add_input_dataset(
+ {"scheme": "s3", "bucket": source_bucket_name, "key": source_bucket_key}, self
+ )
+ get_hook_lineage_collector().add_output_dataset(
+ {"scheme": "s3", "bucket": dest_bucket_name, "key": dest_bucket_key}, self
+ )
return response
@provide_bucket_name
@@ -1425,6 +1441,10 @@ class S3Hook(AwsBaseHook):
file_path.parent.mkdir(exist_ok=True, parents=True)
+ get_hook_lineage_collector().add_output_dataset(
+ {"scheme": "file", "path": file_path if file_path.is_absolute() else file_path.absolute()},
+ self,
+ )
file = open(file_path, "wb")
else:
file = NamedTemporaryFile(dir=local_path, prefix="airflow_tmp_", delete=False) # type: ignore
@@ -1435,7 +1455,9 @@ class S3Hook(AwsBaseHook):
ExtraArgs=self.extra_args,
Config=self.transfer_config,
)
-
+ get_hook_lineage_collector().add_input_dataset(
+ {"scheme": "s3", "bucket": bucket_name, "key": key}, self
+ )
return file.name
def generate_presigned_url(
diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml
index 478aeb230c..a2ead1ea5a 100644
--- a/airflow/providers/amazon/provider.yaml
+++ b/airflow/providers/amazon/provider.yaml
@@ -549,6 +549,7 @@ sensors:
dataset-uris:
- schemes: [s3]
handler: null
+ factory: airflow.providers.amazon.aws.datasets.s3.create_dataset
filesystems:
- airflow.providers.amazon.aws.fs.s3
diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py
index 9e9dd4d573..d76fc8e016 100644
--- a/airflow/providers_manager.py
+++ b/airflow/providers_manager.py
@@ -901,8 +901,7 @@ class ProvidersManager(LoggingMixin, metaclass=Singleton):
factory_path is not None
and (factory := _correctness_check(provider_package, factory_path, provider))
):
- continue
- self._dataset_factories.update((scheme, factory) for scheme in schemes)
+ self._dataset_factories.update((scheme, factory) for scheme in schemes)
def _discover_taskflow_decorators(self) -> None:
for name, info in self._provider_dict.items():
@@ -1301,7 +1300,9 @@ class ProvidersManager(LoggingMixin, metaclass=Singleton):
@property
def dataset_factories(self) -> dict[str, Callable[..., Dataset]]:
+ self.log.error("PRE")
self.initialize_providers_dataset_uri_handlers_and_factories()
+ self.log.error("????")
return self._dataset_factories
@property
diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json
index bc7e5eee3a..585ff9b150 100644
--- a/generated/provider_dependencies.json
+++ b/generated/provider_dependencies.json
@@ -57,6 +57,7 @@
"cross-providers-deps": [
"apache.hive",
"cncf.kubernetes",
+ "common.compat",
"common.sql",
"exasol",
"ftp",
diff --git a/tests/conftest.py b/tests/conftest.py
index 9027391575..6cb74446dc 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1326,6 +1326,16 @@ def airflow_root_path() -> Path:
return Path(airflow.__path__[0]).parent
+@pytest.fixture
+def hook_lineage_collector():
+ from airflow.lineage import hook
+
+ hook._hook_lineage_collector = None
+ hook._hook_lineage_collector = hook.HookLineageCollector()
+ yield hook.get_hook_lineage_collector()
+ hook._hook_lineage_collector = None
+
+
# This constant is set to True if tests are run with Airflow installed from Packages rather than running
# the tests within Airflow sources. While most tests in CI are run using Airflow sources, there are
# also compatibility tests that only use `tests` package and run against installed packages of Airflow in
diff --git a/tests/providers/amazon/aws/datasets/__init__.py b/tests/providers/amazon/aws/datasets/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/tests/providers/amazon/aws/datasets/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/tests/providers/amazon/aws/datasets/test_s3.py b/tests/providers/amazon/aws/datasets/test_s3.py
new file mode 100644
index 0000000000..c7ffe25240
--- /dev/null
+++ b/tests/providers/amazon/aws/datasets/test_s3.py
@@ -0,0 +1,27 @@
+# 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 airflow.datasets import Dataset
+from airflow.providers.amazon.aws.datasets.s3 import create_dataset
+
+
+def test_create_dataset():
+ assert create_dataset(bucket="test-bucket", key="test-path") == Dataset(uri="s3://test-bucket/test-path")
+ assert create_dataset(bucket="test-bucket", key="test-dir/test-path") == Dataset(
+ uri="s3://test-bucket/test-dir/test-path"
+ )
diff --git a/tests/providers/amazon/aws/hooks/test_s3.py b/tests/providers/amazon/aws/hooks/test_s3.py
index 6b10173d3c..51f2ec2c0a 100644
--- a/tests/providers/amazon/aws/hooks/test_s3.py
+++ b/tests/providers/amazon/aws/hooks/test_s3.py
@@ -31,6 +31,7 @@ import pytest
from botocore.exceptions import ClientError
from moto import mock_aws
+from airflow.datasets import Dataset
from airflow.exceptions import AirflowException
from airflow.models import Connection
from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure
@@ -388,6 +389,14 @@ class TestAwsS3Hook:
resource = boto3.resource("s3").Object(s3_bucket, "my_key")
assert resource.get()["Body"].read() == b"Cont\xc3\xa9nt"
+ def test_load_string_exposes_lineage(self, s3_bucket, hook_lineage_collector):
+ hook = S3Hook()
+ hook.load_string("Contént", "my_key", s3_bucket)
+ assert len(hook_lineage_collector.collected_datasets.outputs) == 1
+ assert hook_lineage_collector.collected_datasets.outputs[0][0] == Dataset(
+ uri=f"s3://{s3_bucket}/my_key"
+ )
+
def test_load_string_compress(self, s3_bucket):
hook = S3Hook()
hook.load_string("Contént", "my_key", s3_bucket, compression="gzip")
@@ -970,6 +979,16 @@ class TestAwsS3Hook:
resource = boto3.resource("s3").Object(s3_bucket, "my_key")
assert gz.decompress(resource.get()["Body"].read()) == b"Content"
+ def test_load_file_exposes_lineage(self, s3_bucket, tmp_path, hook_lineage_collector):
+ hook = S3Hook()
+ path = tmp_path / "testfile"
+ path.write_text("Content")
+ hook.load_file(path, "my_key", s3_bucket)
+ assert len(hook_lineage_collector.collected_datasets.outputs) == 1
+ assert hook_lineage_collector.collected_datasets.outputs[0][0] == Dataset(
+ uri=f"s3://{s3_bucket}/my_key"
+ )
+
def test_load_file_acl(self, s3_bucket, tmp_path):
hook = S3Hook()
path = tmp_path / "testfile"
@@ -1027,6 +1046,25 @@ class TestAwsS3Hook:
ACL="private",
)
+ @mock_aws
+ def test_copy_object_ol_instrumentation(self, s3_bucket, hook_lineage_collector):
+ mock_hook = S3Hook()
+
+ with mock.patch.object(
+ S3Hook,
+ "get_conn",
+ ):
+ mock_hook.copy_object("my_key", "my_key3", s3_bucket, s3_bucket)
+ assert len(hook_lineage_collector.collected_datasets.inputs) == 1
+ assert hook_lineage_collector.collected_datasets.inputs[0][0] == Dataset(
+ uri=f"s3://{s3_bucket}/my_key"
+ )
+
+ assert len(hook_lineage_collector.collected_datasets.outputs) == 1
+ assert hook_lineage_collector.collected_datasets.outputs[0][0] == Dataset(
+ uri=f"s3://{s3_bucket}/my_key3"
+ )
+
@mock_aws
def test_delete_bucket_if_bucket_exist(self, s3_bucket):
# assert if the bucket is created
@@ -1118,7 +1156,7 @@ class TestAwsS3Hook:
assert isinstance(ctx.value, ValueError)
@mock.patch("airflow.providers.amazon.aws.hooks.s3.NamedTemporaryFile")
- def test_download_file(self, mock_temp_file, tmp_path):
+ def test_download_file(self, mock_temp_file, tmp_path, hook_lineage_collector):
path = tmp_path / "airflow_tmp_test_s3_hook"
mock_temp_file.return_value = path
s3_hook = S3Hook(aws_conn_id="s3_test")
@@ -1139,9 +1177,13 @@ class TestAwsS3Hook:
)
assert path.name == output_file
+ assert len(hook_lineage_collector.collected_datasets.inputs) == 1
+ assert hook_lineage_collector.collected_datasets.inputs[0][0] == Dataset(
+ uri="s3://test_bucket/test_key"
+ )
@mock.patch("airflow.providers.amazon.aws.hooks.s3.open")
- def test_download_file_with_preserve_name(self, mock_open, tmp_path):
+ def test_download_file_with_preserve_name(self, mock_open, tmp_path, hook_lineage_collector):
path = tmp_path / "test.log"
bucket = "test_bucket"
key = f"test_key/{path.name}"
@@ -1152,15 +1194,25 @@ class TestAwsS3Hook:
s3_obj.key = f"s3://{bucket}/{key}"
s3_obj.download_fileobj = Mock(return_value=None)
s3_hook.get_key = Mock(return_value=s3_obj)
+ local_path = os.fspath(path.parent)
s3_hook.download_file(
key=key,
bucket_name=bucket,
- local_path=os.fspath(path.parent),
+ local_path=local_path,
preserve_file_name=True,
use_autogenerated_subdir=False,
)
mock_open.assert_called_once_with(path, "wb")
+ assert len(hook_lineage_collector.collected_datasets.inputs) == 1
+ assert hook_lineage_collector.collected_datasets.inputs[0][0] == Dataset(
+ uri="s3://test_bucket/test_key/test.log"
+ )
+
+ assert len(hook_lineage_collector.collected_datasets.outputs) == 1
+ assert hook_lineage_collector.collected_datasets.outputs[0][0] == Dataset(
+ uri=f"file://{local_path}/test.log",
+ )
@mock.patch("airflow.providers.amazon.aws.hooks.s3.open")
def test_download_file_with_preserve_name_with_autogenerated_subdir(self, mock_open, tmp_path):