AIP-62: Add HookLineageCollector

airflow/lineage/hook.py

@@ -0,0 +1,181 @@
+#
+# 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 typing import Union
+
+import attr
+
+from airflow.datasets import Dataset
+from airflow.hooks.base import BaseHook
+from airflow.io.store import ObjectStore
+from airflow.providers_manager import ProvidersManager
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+# Store context what sent lineage.
+LineageContext = Union[BaseHook, ObjectStore]
+
+_hook_lineage_collector: HookLineageCollector | None = None
+
+
+@attr.define
+class HookLineage:
+    """Holds lineage collected by HookLineageCollector."""
+
+    inputs: list[tuple[Dataset, LineageContext]] = attr.ib(factory=list)
+    outputs: list[tuple[Dataset, LineageContext]] = attr.ib(factory=list)
+
+
+class HookLineageCollector(LoggingMixin):
+    """
+    HookLineageCollector is a base class for collecting hook lineage information.
+
+    It is used to collect the input and output datasets of a hook execution.
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.inputs: list[tuple[Dataset, LineageContext]] = []
+        self.outputs: list[tuple[Dataset, LineageContext]] = []
+
+    def create_dataset(
+        self, scheme: str | None, uri: str | None, dataset_kwargs: dict | None, dataset_extra: dict | None
+    ) -> Dataset | None:
+        """
+        Create a Dataset instance using the provided parameters.
+
+        This method attempts to create a Dataset instance using the given parameters.
+        It first checks if a URI is provided and falls back to using the default dataset factory
+        with the given URI if no other information is available.
+
+        If a scheme is provided but no URI, it attempts to find a dataset factory that matches
+        the given scheme. If no such factory is found, it logs an error message and returns None.
+
+        If dataset_kwargs is provided, it is used to pass additional parameters to the Dataset
+        factory. The dataset_extra parameter is also passed to the factory as an ``extra`` parameter.
+        """
+        if uri:
+            # Fallback to default factory using the provided URI
+            return Dataset(uri=uri, extra=dataset_extra)
+
+        if not scheme:
+            self.log.debug(
+                "Missing required parameter: either 'uri' or 'scheme' must be provided to create a Dataset."
+            )
+            return None
+
+        dataset_factory = ProvidersManager().dataset_factories.get(scheme)
+        if not dataset_factory:
+            self.log.debug("Unsupported scheme: %s. Please provide a valid URI to create a Dataset.", scheme)
+            return None
+
+        dataset_kwargs = dataset_kwargs or {}
+        try:
+            return dataset_factory(**dataset_kwargs, extra=dataset_extra)
+        except Exception as e:
+            self.log.debug("Failed to create dataset. Skipping. Error: %s", e)
+            return None
+
+    def add_input_dataset(
+        self,
+        context: LineageContext,
+        scheme: str | None = None,
+        uri: str | None = None,
+        dataset_kwargs: dict | None = None,
+        dataset_extra: dict | None = None,
+    ):
+        """Add the input dataset and its corresponding hook execution context to the collector."""
+        dataset = self.create_dataset(
+            scheme=scheme, uri=uri, dataset_kwargs=dataset_kwargs, dataset_extra=dataset_extra
+        )
+        if dataset:
+            self.inputs.append((dataset, context))
+
+    def add_output_dataset(
+        self,
+        context: LineageContext,
+        scheme: str | None = None,
+        uri: str | None = None,
+        dataset_kwargs: dict | None = None,
+        dataset_extra: dict | None = None,
+    ):
+        """Add the output dataset and its corresponding hook execution context to the collector."""
+        dataset = self.create_dataset(
+            scheme=scheme, uri=uri, dataset_kwargs=dataset_kwargs, dataset_extra=dataset_extra
+        )
+        if dataset:
+            self.outputs.append((dataset, context))
+
+    @property
+    def collected_datasets(self) -> HookLineage:
+        """Get the collected hook lineage information."""
+        return HookLineage(self.inputs, self.outputs)
+
+    @property
+    def has_collected(self) -> bool:
+        """Check if any datasets have been collected."""
+        return len(self.inputs) != 0 or len(self.outputs) != 0
+
+
+class NoOpCollector(HookLineageCollector):
+    """
+    NoOpCollector is a hook lineage collector that does nothing.
+
+    It is used when you want to disable lineage collection.
+    """
+
+    def add_input_dataset(self, *_):
+        pass
+
+    def add_output_dataset(self, *_):
+        pass
+
+    @property
+    def collected_datasets(
+        self,
+    ) -> HookLineage:
+        self.log.warning(
+            "Data lineage tracking is disabled. Register a hook lineage reader to start tracking hook lineage."
+        )
+        return HookLineage([], [])
+
+
+class HookLineageReader(LoggingMixin):
+    """Class used to retrieve the hook lineage information collected by HookLineageCollector."""
+
+    def __init__(self, **kwargs):
+        self.lineage_collector = get_hook_lineage_collector()
+
+    def retrieve_hook_lineage(self) -> HookLineage:
+        """Retrieve hook lineage from HookLineageCollector."""
+        hook_lineage = self.lineage_collector.collected_datasets
+        return hook_lineage
+
+
+def get_hook_lineage_collector() -> HookLineageCollector:
+    """Get singleton lineage collector."""
+    global _hook_lineage_collector
+    if not _hook_lineage_collector:
+        from airflow import plugins_manager
+
+        plugins_manager.initialize_hook_lineage_readers_plugins()
+        if plugins_manager.hook_lineage_reader_classes:
+            _hook_lineage_collector = HookLineageCollector()
+        else:
+            _hook_lineage_collector = NoOpCollector()
+    return _hook_lineage_collector

airflow/plugins_manager.py

@@ -27,6 +27,7 @@
 import os
 import sys
 import types
+from cgitb import Hook
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Iterable
 
@@ -41,6 +42,8 @@
 from airflow.utils.module_loading import import_string, qualname
 
 if TYPE_CHECKING:
+    from airflow.lineage.hook import HookLineageReader
+
     try:
         import importlib_metadata as metadata
     except ImportError:
@@ -75,6 +78,7 @@
 registered_operator_link_classes: dict[str, type] | None = None
 registered_ti_dep_classes: dict[str, type] | None = None
 timetable_classes: dict[str, type[Timetable]] | None = None
+hook_lineage_reader_classes: list[type[Hook]] | None = None
 priority_weight_strategy_classes: dict[str, type[PriorityWeightStrategy]] | None = None
 """
 Mapping of class names to class of OperatorLinks registered by plugins.
@@ -176,8 +180,12 @@ class AirflowPlugin:
     # A list of timetable classes that can be used for DAG scheduling.
     timetables: list[type[Timetable]] = []
 
+    # A list of listeners that can be used for tracking task and DAG states.
     listeners: list[ModuleType | object] = []
 
+    # A list of hook lineage reader classes that can be used for reading lineage information from a hook.
+    hook_lineage_readers: list[type[HookLineageReader]] = []
+
     # A list of priority weight strategy classes that can be used for calculating tasks weight priority.
     priority_weight_strategies: list[type[PriorityWeightStrategy]] = []
 
@@ -483,6 +491,25 @@ def initialize_timetables_plugins():
     }
 
 
+def initialize_hook_lineage_readers_plugins():
+    """Collect hook lineage reader classes registered by plugins."""
+    global hook_lineage_reader_classes
+
+    if hook_lineage_reader_classes is not None:
+        return
+
+    ensure_plugins_loaded()
+
+    if plugins is None:
+        raise AirflowPluginException("Can't load plugins.")
+
+    log.debug("Initialize hook lineage readers plugins")
+
+    hook_lineage_reader_classes = []
+    for plugin in plugins:
+        hook_lineage_reader_classes.extend(plugin.hook_lineage_readers)
+
+
 def integrate_executor_plugins() -> None:
     """Integrate executor plugins to the context."""
     global plugins

airflow/provider.yaml.schema.json

@@ -212,6 +212,10 @@
                     "handler": {
                         "type": ["string", "null"],
                         "description": "Normalization function for specified URI schemes. Import path to a callable taking and returning a SplitResult. 'null' specifies a no-op."
+                    },
+                    "factory": {
+                        "type": ["string", "null"],
+                        "description": "Dataset factory for specified URI. Creates AIP-60 compliant Dataset."
                     }
                 }
             }

airflow/providers_manager.py

@@ -91,6 +91,7 @@ def ensure_prefix(field):
 if TYPE_CHECKING:
     from urllib.parse import SplitResult
 
+    from airflow.datasets import Dataset
     from airflow.decorators.base import TaskDecorator
     from airflow.hooks.base import BaseHook
     from airflow.typing_compat import Literal
@@ -426,6 +427,7 @@ def __init__(self):
         self._hooks_dict: dict[str, HookInfo] = {}
         self._fs_set: set[str] = set()
         self._dataset_uri_handlers: dict[str, Callable[[SplitResult], SplitResult]] = {}
+        self._dataset_factories: dict[str, Callable[..., Dataset]] = {}
         self._taskflow_decorators: dict[str, Callable] = LazyDictWithCache()  # type: ignore[assignment]
         # keeps mapping between connection_types and hook class, package they come from
         self._hook_provider_dict: dict[str, HookClassProvider] = {}
@@ -523,11 +525,12 @@ def initialize_providers_filesystems(self):
         self._discover_filesystems()
 
     @provider_info_cache("dataset_uris")
-    def initialize_providers_dataset_uri_handlers(self):
+    def initialize_providers_dataset_uri_handlers_and_factories(self):
         """Lazy initialization of provider dataset URI handlers."""
         self.initialize_providers_list()
-        self._discover_dataset_uri_handlers()
+        self._discover_dataset_uri_handlers_and_factories()
 
+    @provider_info_cache("hook_lineage_writers")
     @provider_info_cache("taskflow_decorators")
     def initialize_providers_taskflow_decorator(self):
         """Lazy initialization of providers hooks."""
@@ -878,7 +881,7 @@ def _discover_filesystems(self) -> None:
                     self._fs_set.add(fs_module_name)
         self._fs_set = set(sorted(self._fs_set))
 
-    def _discover_dataset_uri_handlers(self) -> None:
+    def _discover_dataset_uri_handlers_and_factories(self) -> None:
         from airflow.datasets import normalize_noop
 
         for provider_package, provider in self._provider_dict.items():
@@ -893,6 +896,13 @@ def _discover_dataset_uri_handlers(self) -> None:
                 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))
+                ):
+                    continue
+                self._dataset_factories.update((scheme, factory) for scheme in schemes)
 
     def _discover_taskflow_decorators(self) -> None:
         for name, info in self._provider_dict.items():
@@ -1289,9 +1299,14 @@ def filesystem_module_names(self) -> list[str]:
         self.initialize_providers_filesystems()
         return sorted(self._fs_set)
 
+    @property
+    def dataset_factories(self) -> dict[str, Callable[..., Dataset]]:
+        self.initialize_providers_dataset_uri_handlers_and_factories()
+        return self._dataset_factories
+
     @property
     def dataset_uri_handlers(self) -> dict[str, Callable[[SplitResult], SplitResult]]:
-        self.initialize_providers_dataset_uri_handlers()
+        self.initialize_providers_dataset_uri_handlers_and_factories()
         return self._dataset_uri_handlers
 
     @property

docs/apache-airflow/administration-and-deployment/lineage.rst

@@ -89,6 +89,49 @@ has outlets defined (e.g. by using ``add_outlets(..)`` or has out of the box sup
 
 .. _precedence: https://docs.python.org/3/reference/expressions.html
 
+Hook Lineage
+------------
+
+Airflow provides a powerful feature for tracking data lineage not only between tasks but also from hooks used within those tasks.
+This functionality helps you understand how data flows throughout your Airflow pipelines.
+
+A global instance of ``HookLineageCollector`` serves as the central hub for collecting lineage information.
+Hooks can send details about datasets they interact with to this collector.
+The collector then uses this data to construct AIP-60 compliant Datasets, a standard format for describing datasets.
+
+.. code-block:: python
+
+    from airflow.lineage.hook_lineage import get_hook_lineage_collector
+
+
+    class CustomHook(BaseHook):
+        def run(self):
+            # run actual code
+            collector = get_hook_lineage_collector()
+            collector.add_input_dataset(self, dataset_kwargs={"scheme": "file", "path": "/tmp/in"})
+            collector.add_output_dataset(self, dataset_kwargs={"scheme": "file", "path": "/tmp/out"})
+
+Lineage data collected by the ``HookLineageCollector`` can be accessed using an instance of ``HookLineageReader``,
+which is registered in an Airflow plugin.
+
+.. code-block:: python
+
+    from airflow.lineage.hook_lineage import HookLineageReader
+    from airflow.plugins_manager import AirflowPlugin
+
+
+    class CustomHookLineageReader(HookLineageReader):
+        def get_inputs(self):
+            return self.lineage_collector.collected_datasets.inputs
+
+
+    class HookLineageCollectionPlugin(AirflowPlugin):
+        name = "HookLineageCollectionPlugin"
+        hook_lineage_readers = [CustomHookLineageReader]
+
+If no ``HookLineageReader`` is registered within Airflow, a default ``NoOpCollector`` is used instead.
+This collector does not create AIP-60 compliant datasets or collect lineage information.
+
 
 Lineage Backend
 ---------------