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:46:34 UTC
(airflow) branch aip-62/s3hook created (now ba1bb1268a)
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
at ba1bb1268a openlineage: add support for hook lineage for S3Hook
This branch includes the following new commits:
new ba1bb1268a openlineage: add support for hook lineage for S3Hook
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.
(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 ba1bb1268a8157dcf614c1e0faa0b4eca106e4fd
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/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/hooks/test_s3.py | 58 +++++++++++++++++++++++++++--
6 files changed, 93 insertions(+), 6 deletions(-)
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/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):