openlineage: add support for hook lineage for S3Hook (#40819)

mobuchowski · web-flow · commit 0212f671921f · 2024-07-22T11:00:48.000+02:00
Signed-off-by: Maciej Obuchowski &lt;obuchowski.maciej@gmail.com&gt;
diff --git a/airflow/lineage/hook.py b/airflow/lineage/hook.py
@@ -139,10 +139,10 @@ class NoOpCollector(HookLineageCollector):
     It is used when you want to disable lineage collection.
     """
 
-    def add_input_dataset(self, *_):
+    def add_input_dataset(self, *_, **__):
         pass
 
-    def add_output_dataset(self, *_):
+    def add_output_dataset(self, *_, **__):
         pass
 
     @property
diff --git a/airflow/providers/amazon/aws/datasets/__init__.py 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
@@ -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, extra=None) -> Dataset:
+    return Dataset(uri=f"s3://{bucket}/{key}", extra=extra)
diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
@@ -41,6 +41,8 @@
 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,12 @@ def load_file(
 
         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(
+            context=self, scheme="file", dataset_kwargs={"path": filename}
+        )
+        get_hook_lineage_collector().add_output_dataset(
+            context=self, scheme="s3", dataset_kwargs={"bucket": bucket_name, "key": key}
+        )
 
     @unify_bucket_name_and_key
     @provide_bucket_name
@@ -1251,6 +1259,10 @@ def _upload_file_obj(
             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(
+            context=self, scheme="s3", dataset_kwargs={"bucket": bucket_name, "key": key}
+        )
 
     def copy_object(
         self,
@@ -1306,6 +1318,12 @@ def copy_object(
         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(
+            context=self, scheme="s3", dataset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}
+        )
+        get_hook_lineage_collector().add_output_dataset(
+            context=self, scheme="s3", dataset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}
+        )
         return response
 
     @provide_bucket_name
@@ -1425,6 +1443,11 @@ def download_file(
 
             file_path.parent.mkdir(exist_ok=True, parents=True)
 
+            get_hook_lineage_collector().add_output_dataset(
+                context=self,
+                scheme="file",
+                dataset_kwargs={"path": file_path if file_path.is_absolute() else file_path.absolute()},
+            )
             file = open(file_path, "wb")
         else:
             file = NamedTemporaryFile(dir=local_path, prefix="airflow_tmp_", delete=False)  # type: ignore
@@ -1435,7 +1458,9 @@ def download_file(
                 ExtraArgs=self.extra_args,
                 Config=self.transfer_config,
             )
-
+        get_hook_lineage_collector().add_input_dataset(
+            context=self, scheme="s3", dataset_kwargs={"bucket": bucket_name, "key": key}
+        )
         return file.name
 
     def generate_presigned_url(
diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml
@@ -91,6 +91,7 @@ dependencies:
   - apache-airflow>=2.7.0
   - apache-airflow-providers-common-sql>=1.3.1
   - apache-airflow-providers-http
+  - apache-airflow-providers-common-compat>=1.1.0
   # We should update minimum version of boto3 and here regularly to avoid `pip` backtracking with the number
   # of candidates to consider. Make sure to configure boto3 version here as well as in all the tools below
   # in the `devel-dependencies` section to be the same minimum version.
@@ -561,6 +562,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/common/compat/lineage/hook.py b/airflow/providers/common/compat/lineage/hook.py
@@ -32,10 +32,10 @@ class NoOpCollector:
             It is used when you want to disable lineage collection.
             """
 
-            def add_input_dataset(self, *_):
+            def add_input_dataset(self, *_, **__):
                 pass
 
-            def add_output_dataset(self, *_):
+            def add_output_dataset(self, *_, **__):
                 pass
 
         return NoOpCollector()
diff --git a/airflow/providers/common/compat/provider.yaml b/airflow/providers/common/compat/provider.yaml
@@ -25,6 +25,7 @@ state: ready
 source-date-epoch: 1716287191
 # note that those versions are maintained by release manager - do not update them manually
 versions:
+  - 1.1.0
   - 1.0.0
 
 dependencies:
diff --git a/airflow/providers/common/io/datasets/file.py b/airflow/providers/common/io/datasets/file.py
@@ -19,6 +19,6 @@
 from airflow.datasets import Dataset
 
 
-def create_dataset(*, path: str) -> Dataset:
+def create_dataset(*, path: str, extra=None) -> Dataset:
     # We assume that we get absolute path starting with /
-    return Dataset(uri=f"file://{path}")
+    return Dataset(uri=f"file://{path}", extra=extra)
diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py
@@ -886,23 +886,23 @@ def _discover_dataset_uri_handlers_and_factories(self) -> None:
 
         for provider_package, provider in self._provider_dict.items():
             for handler_info in provider.data.get("dataset-uris", []):
-                try:
-                    schemes = handler_info["schemes"]
-                    handler_path = handler_info["handler"]
-                except KeyError:
+                schemes = handler_info.get("schemes")
+                handler_path = handler_info.get("handler")
+                factory_path = handler_info.get("factory")
+                if schemes is None:
                     continue
-                if handler_path is None:
+
+                if handler_path is not None and (
+                    handler := _correctness_check(provider_package, handler_path, provider)
+                ):
+                    pass
+                else:
                     handler = normalize_noop
-                elif not (handler := _correctness_check(provider_package, handler_path, provider)):
-                    continue
                 self._dataset_uri_handlers.update((scheme, handler) for scheme in schemes)
-                factory_path = handler_info.get("factory")
-                if not (
-                    factory_path is not None
-                    and (factory := _correctness_check(provider_package, factory_path, provider))
+                if 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():
diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py
@@ -569,7 +569,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
             ("airflow/providers/amazon/__init__.py",),
             {
                 "affected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
-                "common.sql exasol ftp google http imap microsoft.azure "
+                "common.compat common.sql exasol ftp google http imap microsoft.azure "
                 "mongo mysql openlineage postgres salesforce ssh teradata",
                 "all-python-versions": "['3.8']",
                 "all-python-versions-list-as-string": "3.8",
@@ -585,7 +585,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
                 "upgrade-to-newer-dependencies": "false",
                 "run-amazon-tests": "true",
                 "parallel-test-types-list-as-string": "Always Providers[amazon] "
-                "Providers[apache.hive,cncf.kubernetes,common.sql,exasol,ftp,http,"
+                "Providers[apache.hive,cncf.kubernetes,common.compat,common.sql,exasol,ftp,http,"
                 "imap,microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] Providers[google]",
                 "needs-mypy": "true",
                 "mypy-folders": "['providers']",
@@ -619,7 +619,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
             ("airflow/providers/amazon/file.py",),
             {
                 "affected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
-                "common.sql exasol ftp google http imap microsoft.azure "
+                "common.compat common.sql exasol ftp google http imap microsoft.azure "
                 "mongo mysql openlineage postgres salesforce ssh teradata",
                 "all-python-versions": "['3.8']",
                 "all-python-versions-list-as-string": "3.8",
@@ -635,7 +635,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
                 "run-kubernetes-tests": "false",
                 "upgrade-to-newer-dependencies": "false",
                 "parallel-test-types-list-as-string": "Always Providers[amazon] "
-                "Providers[apache.hive,cncf.kubernetes,common.sql,exasol,ftp,http,"
+                "Providers[apache.hive,cncf.kubernetes,common.compat,common.sql,exasol,ftp,http,"
                 "imap,microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] Providers[google]",
                 "needs-mypy": "true",
                 "mypy-folders": "['providers']",
diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json
@@ -28,6 +28,7 @@
   "amazon": {
     "deps": [
       "PyAthena>=3.0.10",
+      "apache-airflow-providers-common-compat>=1.1.0",
       "apache-airflow-providers-common-sql>=1.3.1",
       "apache-airflow-providers-http",
       "apache-airflow>=2.7.0",
@@ -57,6 +58,7 @@
     "cross-providers-deps": [
       "apache.hive",
       "cncf.kubernetes",
+      "common.compat",
       "common.sql",
       "exasol",
       "ftp",
diff --git a/prod_image_installed_providers.txt b/prod_image_installed_providers.txt
@@ -2,6 +2,7 @@
 amazon
 celery
 cncf.kubernetes
+common.compat
 common.io
 common.sql
 docker
diff --git 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
@@ -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
@@ -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
