apache-airflow-providers-openlineage 1.3.1rc1__py3-none-any.whl

airflow/providers/openlineage/extractors/manager.py

@@ -0,0 +1,199 @@
+# 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
+
+import os
+from contextlib import suppress
+from typing import TYPE_CHECKING, Iterator
+
+from airflow.configuration import conf
+from airflow.providers.openlineage.extractors import BaseExtractor, OperatorLineage
+from airflow.providers.openlineage.extractors.base import DefaultExtractor
+from airflow.providers.openlineage.extractors.bash import BashExtractor
+from airflow.providers.openlineage.extractors.python import PythonExtractor
+from airflow.providers.openlineage.plugins.facets import (
+    UnknownOperatorAttributeRunFacet,
+    UnknownOperatorInstance,
+)
+from airflow.providers.openlineage.utils.utils import get_filtered_unknown_operator_keys
+from airflow.utils.log.logging_mixin import LoggingMixin
+from airflow.utils.module_loading import import_string
+
+if TYPE_CHECKING:
+    from airflow.models import Operator
+
+
+def try_import_from_string(string):
+    with suppress(ImportError):
+        return import_string(string)
+
+
+def _iter_extractor_types() -> Iterator[type[BaseExtractor]]:
+    if PythonExtractor is not None:
+        yield PythonExtractor
+    if BashExtractor is not None:
+        yield BashExtractor
+
+
+class ExtractorManager(LoggingMixin):
+    """Class abstracting management of custom extractors."""
+
+    def __init__(self):
+        super().__init__()
+        self.extractors: dict[str, type[BaseExtractor]] = {}
+        self.default_extractor = DefaultExtractor
+
+        # Comma-separated extractors in OPENLINEAGE_EXTRACTORS variable.
+        # Extractors should implement BaseExtractor
+        for extractor in _iter_extractor_types():
+            for operator_class in extractor.get_operator_classnames():
+                self.extractors[operator_class] = extractor
+
+        env_extractors = conf.get("openlineage", "extractors", fallback=os.getenv("OPENLINEAGE_EXTRACTORS"))
+        # skip either when it's empty string or None
+        if env_extractors:
+            for extractor in env_extractors.split(";"):
+                extractor: type[BaseExtractor] = try_import_from_string(extractor.strip())
+                for operator_class in extractor.get_operator_classnames():
+                    self.extractors[operator_class] = extractor
+
+    def add_extractor(self, operator_class: str, extractor: type[BaseExtractor]):
+        self.extractors[operator_class] = extractor
+
+    def extract_metadata(self, dagrun, task, complete: bool = False, task_instance=None) -> OperatorLineage:
+        extractor = self._get_extractor(task)
+        task_info = (
+            f"task_type={task.task_type} "
+            f"airflow_dag_id={task.dag_id} "
+            f"task_id={task.task_id} "
+            f"airflow_run_id={dagrun.run_id} "
+        )
+
+        if extractor:
+            # Extracting advanced metadata is only possible when extractor for particular operator
+            # is defined. Without it, we can't extract any input or output data.
+            try:
+                self.log.debug("Using extractor %s %s", extractor.__class__.__name__, str(task_info))
+                if complete:
+                    task_metadata = extractor.extract_on_complete(task_instance)
+                else:
+                    task_metadata = extractor.extract()
+
+                self.log.debug("Found task metadata for operation %s: %s", task.task_id, str(task_metadata))
+                task_metadata = self.validate_task_metadata(task_metadata)
+                if task_metadata:
+                    if (not task_metadata.inputs) and (not task_metadata.outputs):
+                        self.extract_inlets_and_outlets(task_metadata, task.inlets, task.outlets)
+
+                    return task_metadata
+
+            except Exception as e:
+                self.log.warning(
+                    "Failed to extract metadata using found extractor %s - %s %s", extractor, e, task_info
+                )
+        else:
+            self.log.debug("Unable to find an extractor %s", task_info)
+
+            # Only include the unkonwnSourceAttribute facet if there is no extractor
+            task_metadata = OperatorLineage(
+                run_facets={
+                    "unknownSourceAttribute": UnknownOperatorAttributeRunFacet(
+                        unknownItems=[
+                            UnknownOperatorInstance(
+                                name=task.task_type,
+                                properties=get_filtered_unknown_operator_keys(task),
+                            )
+                        ]
+                    )
+                },
+            )
+            inlets = task.get_inlet_defs()
+            outlets = task.get_outlet_defs()
+            self.extract_inlets_and_outlets(task_metadata, inlets, outlets)
+            return task_metadata
+
+        return OperatorLineage()
+
+    def get_extractor_class(self, task: Operator) -> type[BaseExtractor] | None:
+        if task.task_type in self.extractors:
+            return self.extractors[task.task_type]
+
+        def method_exists(method_name):
+            method = getattr(task, method_name, None)
+            if method:
+                return callable(method)
+
+        if method_exists("get_openlineage_facets_on_start") or method_exists(
+            "get_openlineage_facets_on_complete"
+        ):
+            return self.default_extractor
+        return None
+
+    def _get_extractor(self, task: Operator) -> BaseExtractor | None:
+        # TODO: Re-enable in Extractor PR
+        # self.instantiate_abstract_extractors(task)
+        extractor = self.get_extractor_class(task)
+        self.log.debug("extractor for %s is %s", task.task_type, extractor)
+        if extractor:
+            return extractor(task)
+        return None
+
+    def extract_inlets_and_outlets(
+        self,
+        task_metadata: OperatorLineage,
+        inlets: list,
+        outlets: list,
+    ):
+        if inlets or outlets:
+            self.log.debug("Manually extracting lineage metadata from inlets and outlets")
+        for i in inlets:
+            d = self.convert_to_ol_dataset(i)
+            if d:
+                task_metadata.inputs.append(d)
+        for o in outlets:
+            d = self.convert_to_ol_dataset(o)
+            if d:
+                task_metadata.outputs.append(d)
+
+    @staticmethod
+    def convert_to_ol_dataset(obj):
+        from openlineage.client.run import Dataset
+
+        from airflow.lineage.entities import Table
+
+        if isinstance(obj, Dataset):
+            return obj
+        elif isinstance(obj, Table):
+            return Dataset(
+                namespace=f"{obj.cluster}",
+                name=f"{obj.database}.{obj.name}",
+                facets={},
+            )
+        else:
+            return None
+
+    def validate_task_metadata(self, task_metadata) -> OperatorLineage | None:
+        try:
+            return OperatorLineage(
+                inputs=task_metadata.inputs,
+                outputs=task_metadata.outputs,
+                run_facets=task_metadata.run_facets,
+                job_facets=task_metadata.job_facets,
+            )
+        except AttributeError:
+            self.log.warning("Extractor returns non-valid metadata: %s", task_metadata)
+            return None

airflow/providers/openlineage/extractors/python.py

@@ -0,0 +1,97 @@
+# 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
+
+import inspect
+from typing import Callable
+
+from openlineage.client.facet import SourceCodeJobFacet
+
+from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
+from airflow.providers.openlineage.plugins.facets import (
+    UnknownOperatorAttributeRunFacet,
+    UnknownOperatorInstance,
+)
+from airflow.providers.openlineage.utils.utils import (
+    get_filtered_unknown_operator_keys,
+    is_source_enabled,
+)
+
+"""
+:meta private:
+"""
+
+
+class PythonExtractor(BaseExtractor):
+    """
+    Extract executed source code and put it into SourceCodeJobFacet.
+
+    This extractor provides visibility on what particular task does by extracting
+    executed source code and putting it into SourceCodeJobFacet. It does not extract
+    datasets yet.
+
+    :meta private:
+    """
+
+    @classmethod
+    def get_operator_classnames(cls) -> list[str]:
+        return ["PythonOperator"]
+
+    def _execute_extraction(self) -> OperatorLineage | None:
+        source_code = self.get_source_code(self.operator.python_callable)
+        job_facet: dict = {}
+        if is_source_enabled() and source_code:
+            job_facet = {
+                "sourceCode": SourceCodeJobFacet(
+                    language="python",
+                    # We're on worker and should have access to DAG files
+                    source=source_code,
+                )
+            }
+        return OperatorLineage(
+            job_facets=job_facet,
+            run_facets={
+                # The PythonOperator is recorded as an "unknownSource" even though we have an
+                # extractor, as the data lineage cannot be determined from the operator
+                # directly.
+                "unknownSourceAttribute": UnknownOperatorAttributeRunFacet(
+                    unknownItems=[
+                        UnknownOperatorInstance(
+                            name="PythonOperator",
+                            properties=get_filtered_unknown_operator_keys(self.operator),
+                        )
+                    ]
+                )
+            },
+        )
+
+    def get_source_code(self, callable: Callable) -> str | None:
+        try:
+            return inspect.getsource(callable)
+        except TypeError:
+            # Trying to extract source code of builtin_function_or_method
+            return str(callable)
+        except OSError:
+            self.log.warning(
+                "Can't get source code facet of PythonOperator %s",
+                self.operator.task_id,
+            )
+        return None
+
+    def extract(self) -> OperatorLineage | None:
+        return super().extract()

airflow/providers/openlineage/get_provider_info.py

@@ -0,0 +1,110 @@
+# 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.
+
+# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
+# OVERWRITTEN WHEN PREPARING PACKAGES.
+#
+# IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
+# `get_provider_info_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
+
+
+def get_provider_info():
+    return {
+        "package-name": "apache-airflow-providers-openlineage",
+        "name": "OpenLineage Airflow",
+        "description": "`OpenLineage <https://openlineage.io/>`__\n",
+        "suspended": False,
+        "source-date-epoch": 1703288155,
+        "versions": ["1.3.1", "1.3.0", "1.2.1", "1.2.0", "1.1.1", "1.1.0", "1.0.2", "1.0.1", "1.0.0"],
+        "dependencies": [
+            "apache-airflow>=2.7.0",
+            "apache-airflow-providers-common-sql>=1.6.0",
+            "attrs>=22.2",
+            "openlineage-integration-common>=0.28.0",
+            "openlineage-python>=0.28.0",
+        ],
+        "integrations": [
+            {
+                "integration-name": "OpenLineage",
+                "external-doc-url": "https://openlineage.io",
+                "logo": "/integration-logos/openlineage/openlineage.svg",
+                "tags": ["protocol"],
+            }
+        ],
+        "plugins": [
+            {
+                "name": "openlineage",
+                "plugin-class": "airflow.providers.openlineage.plugins.openlineage.OpenLineageProviderPlugin",
+            }
+        ],
+        "config": {
+            "openlineage": {
+                "description": "This section applies settings for OpenLineage integration.\nFor backwards compatibility with `openlineage-python` one can still use\n`openlineage.yml` file or `OPENLINEAGE_` environment variables. However, below\nconfiguration takes precedence over those.\nMore in documentation - https://openlineage.io/docs/client/python#configuration.\n",
+                "options": {
+                    "disabled": {
+                        "description": "Set this to true if you don't want OpenLineage to emit events.\n",
+                        "type": "boolean",
+                        "example": None,
+                        "default": "False",
+                        "version_added": None,
+                    },
+                    "disabled_for_operators": {
+                        "description": "Semicolon separated string of Airflow Operator names to disable\n",
+                        "type": "string",
+                        "example": "airflow.operators.bash.BashOperator;airflow.operators.python.PythonOperator",
+                        "default": "",
+                        "version_added": "1.1.0",
+                    },
+                    "namespace": {
+                        "description": "OpenLineage namespace\n",
+                        "version_added": None,
+                        "type": "string",
+                        "example": "food_delivery",
+                        "default": None,
+                    },
+                    "extractors": {
+                        "description": "Semicolon separated paths to custom OpenLineage extractors.\n",
+                        "type": "string",
+                        "example": "full.path.to.ExtractorClass;full.path.to.AnotherExtractorClass",
+                        "default": None,
+                        "version_added": None,
+                    },
+                    "config_path": {
+                        "description": "Path to YAML config. This provides backwards compatibility to pass config as\n`openlineage.yml` file.\n",
+                        "version_added": None,
+                        "type": "string",
+                        "example": None,
+                        "default": "",
+                    },
+                    "transport": {
+                        "description": "OpenLineage Client transport configuration. It should contain type\nand additional options per each type.\n\nCurrently supported types are:\n\n  * HTTP\n  * Kafka\n  * Console\n",
+                        "type": "string",
+                        "example": '{"type": "http", "url": "http://localhost:5000"}',
+                        "default": "",
+                        "version_added": None,
+                    },
+                    "disable_source_code": {
+                        "description": "If disabled, OpenLineage events do not contain source code of particular\noperators, like PythonOperator.\n",
+                        "default": None,
+                        "example": None,
+                        "type": "boolean",
+                        "version_added": None,
+                    },
+                },
+            }
+        },
+    }

airflow/providers/openlineage/plugins/__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.