apache-airflow-providers-openlineage 1.9.1__py3-none-any.whl → 1.10.0__py3-none-any.whl

airflow/providers/openlineage/utils/utils.py

@@ -24,7 +24,7 @@ import re
 from contextlib import redirect_stdout, suppress
 from functools import wraps
 from io import StringIO
-from typing import TYPE_CHECKING, Any, Iterable
+from typing import TYPE_CHECKING, Any, Callable, Iterable
 
 import attrs
 from deprecated import deprecated
@@ -32,15 +32,16 @@ from openlineage.client.utils import RedactMixin
 from packaging.version import Version
 
 from airflow import __version__ as AIRFLOW_VERSION
+from airflow.datasets import Dataset
 from airflow.exceptions import AirflowProviderDeprecationWarning  # TODO: move this maybe to Airflow's logic?
 from airflow.models import DAG, BaseOperator, MappedOperator
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.plugins.facets import (
+    AirflowDagRunFacet,
     AirflowJobFacet,
     AirflowMappedTaskRunFacet,
     AirflowRunFacet,
     AirflowStateRunFacet,
-    BaseFacet,
     UnknownOperatorAttributeRunFacet,
     UnknownOperatorInstance,
 )
@@ -54,12 +55,16 @@ from airflow.utils.log.secrets_masker import Redactable, Redacted, SecretsMasker
 from airflow.utils.module_loading import import_string
 
 if TYPE_CHECKING:
+    from openlineage.client.event_v2 import Dataset as OpenLineageDataset
+    from openlineage.client.facet_v2 import RunFacet
+
     from airflow.models import DagRun, TaskInstance
+    from airflow.utils.state import TaskInstanceState
 
 
 log = logging.getLogger(__name__)
 _NOMINAL_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
-_IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
+IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
 
 
 def try_import_from_string(string: str) -> Any:
@@ -77,12 +82,52 @@ def get_job_name(task: TaskInstance) -> str:
     return f"{task.dag_id}.{task.task_id}"
 
 
-def get_custom_facets(task_instance: TaskInstance | None = None) -> dict[str, Any]:
-    custom_facets = {}
+def get_airflow_mapped_task_facet(task_instance: TaskInstance) -> dict[str, Any]:
     # check for -1 comes from SmartSensor compatibility with dynamic task mapping
     # this comes from Airflow code
     if hasattr(task_instance, "map_index") and getattr(task_instance, "map_index") != -1:
-        custom_facets["airflow_mappedTask"] = AirflowMappedTaskRunFacet.from_task_instance(task_instance)
+        return {"airflow_mappedTask": AirflowMappedTaskRunFacet.from_task_instance(task_instance)}
+    return {}
+
+
+def get_user_provided_run_facets(ti: TaskInstance, ti_state: TaskInstanceState) -> dict[str, RunFacet]:
+    custom_facets = {}
+
+    # Append custom run facets by executing the custom_run_facet functions.
+    for custom_facet_func in conf.custom_run_facets():
+        try:
+            func: Callable[[TaskInstance, TaskInstanceState], dict[str, RunFacet]] | None = (
+                try_import_from_string(custom_facet_func)
+            )
+            if not func:
+                log.warning(
+                    "OpenLineage is unable to import custom facet function `%s`; will ignore it.",
+                    custom_facet_func,
+                )
+                continue
+            facets: dict[str, RunFacet] | None = func(ti, ti_state)
+            if facets and isinstance(facets, dict):
+                duplicate_facet_keys = [facet_key for facet_key in facets if facet_key in custom_facets]
+                if duplicate_facet_keys:
+                    log.warning(
+                        "Duplicate OpenLineage custom facets key(s) found: `%s` from function `%s`; "
+                        "this will overwrite the previous value.",
+                        ", ".join(duplicate_facet_keys),
+                        custom_facet_func,
+                    )
+                log.debug(
+                    "Adding OpenLineage custom facet with key(s): `%s` from function `%s`.",
+                    tuple(facets),
+                    custom_facet_func,
+                )
+                custom_facets.update(facets)
+        except Exception as exc:
+            log.warning(
+                "Error processing custom facet function `%s`; will ignore it. Error was: %s: %s",
+                custom_facet_func,
+                type(exc).__name__,
+                exc,
+            )
     return custom_facets
 
 
@@ -174,7 +219,19 @@ class InfoJsonEncodable(dict):
                     setattr(self, field, getattr(self.obj, field))
                     self._fields.append(field)
         else:
-            for field, val in self.obj.__dict__.items():
+            if hasattr(self.obj, "__dict__"):
+                obj_fields = self.obj.__dict__
+            elif attrs.has(self.obj.__class__):  # e.g. attrs.define class with slots=True has no __dict__
+                obj_fields = {
+                    field.name: getattr(self.obj, field.name) for field in attrs.fields(self.obj.__class__)
+                }
+            else:
+                raise ValueError(
+                    "Cannot iterate over fields: "
+                    f"The object of type {type(self.obj).__name__} neither has a __dict__ attribute "
+                    "nor is defined as an attrs class."
+                )
+            for field, val in obj_fields.items():
                 if field not in self._fields and field not in self.excludes and field not in self.renames:
                     setattr(self, field, val)
                     self._fields.append(field)
@@ -206,7 +263,7 @@ class DagRunInfo(InfoJsonEncodable):
 class TaskInstanceInfo(InfoJsonEncodable):
     """Defines encoding TaskInstance object to JSON."""
 
-    includes = ["duration", "try_number", "pool", "queued_dttm"]
+    includes = ["duration", "try_number", "pool", "queued_dttm", "log_url"]
     casts = {
         "map_index": lambda ti: (
             ti.map_index if hasattr(ti, "map_index") and getattr(ti, "map_index") != -1 else None
@@ -231,6 +288,7 @@ class TaskInfo(InfoJsonEncodable):
         "_is_teardown": "is_teardown",
     }
     includes = [
+        "deferrable",
         "depends_on_past",
         "downstream_task_ids",
         "execution_timeout",
@@ -259,16 +317,36 @@ class TaskInfo(InfoJsonEncodable):
     ]
     casts = {
         "operator_class": lambda task: task.task_type,
+        "operator_class_path": lambda task: get_fully_qualified_class_name(task),
         "task_group": lambda task: (
             TaskGroupInfo(task.task_group)
             if hasattr(task, "task_group") and getattr(task.task_group, "_group_id", None)
             else None
         ),
-        "inlets": lambda task: [DatasetInfo(inlet) for inlet in task.inlets],
-        "outlets": lambda task: [DatasetInfo(outlet) for outlet in task.outlets],
+        "inlets": lambda task: [DatasetInfo(i) for i in task.inlets if isinstance(i, Dataset)],
+        "outlets": lambda task: [DatasetInfo(o) for o in task.outlets if isinstance(o, Dataset)],
     }
 
 
+class TaskInfoComplete(TaskInfo):
+    """Defines encoding BaseOperator/AbstractOperator object to JSON used when user enables full task info."""
+
+    includes = []
+    excludes = [
+        "_BaseOperator__instantiated",
+        "_dag",
+        "_hook",
+        "_log",
+        "_outlets",
+        "_inlets",
+        "_lock_for_execution",
+        "handler",
+        "params",
+        "python_callable",
+        "retry_delay",
+    ]
+
+
 class TaskGroupInfo(InfoJsonEncodable):
     """Defines encoding TaskGroup object to JSON."""
 
@@ -285,25 +363,36 @@ class TaskGroupInfo(InfoJsonEncodable):
     ]
 
 
+def get_airflow_dag_run_facet(dag_run: DagRun) -> dict[str, RunFacet]:
+    if not dag_run.dag:
+        return {}
+    return {
+        "airflowDagRun": AirflowDagRunFacet(
+            dag=DagInfo(dag_run.dag),
+            dagRun=DagRunInfo(dag_run),
+        )
+    }
+
+
 def get_airflow_run_facet(
     dag_run: DagRun,
     dag: DAG,
     task_instance: TaskInstance,
     task: BaseOperator,
     task_uuid: str,
-) -> dict[str, BaseFacet]:
+) -> dict[str, AirflowRunFacet]:
     return {
         "airflow": AirflowRunFacet(
             dag=DagInfo(dag),
             dagRun=DagRunInfo(dag_run),
             taskInstance=TaskInstanceInfo(task_instance),
-            task=TaskInfo(task),
+            task=TaskInfoComplete(task) if conf.include_full_task_info() else TaskInfo(task),
             taskUuid=task_uuid,
         )
     }
 
 
-def get_airflow_job_facet(dag_run: DagRun) -> dict[str, BaseFacet]:
+def get_airflow_job_facet(dag_run: DagRun) -> dict[str, AirflowJobFacet]:
     if not dag_run.dag:
         return {}
     return {
@@ -315,7 +404,7 @@ def get_airflow_job_facet(dag_run: DagRun) -> dict[str, BaseFacet]:
     }
 
 
-def get_airflow_state_run_facet(dag_run: DagRun) -> dict[str, BaseFacet]:
+def get_airflow_state_run_facet(dag_run: DagRun) -> dict[str, AirflowStateRunFacet]:
     return {
         "airflowState": AirflowStateRunFacet(
             dagRunState=dag_run.get_state(),
@@ -574,6 +663,34 @@ def normalize_sql(sql: str | Iterable[str]):
 
 def should_use_external_connection(hook) -> bool:
     # If we're at Airflow 2.10, the execution is process-isolated, so we can safely run those again.
-    if not _IS_AIRFLOW_2_10_OR_HIGHER:
+    if not IS_AIRFLOW_2_10_OR_HIGHER:
         return hook.__class__.__name__ not in ["SnowflakeHook", "SnowflakeSqlApiHook", "RedshiftSQLHook"]
     return True
+
+
+def translate_airflow_dataset(dataset: Dataset, lineage_context) -> OpenLineageDataset | None:
+    """
+    Convert a Dataset with an AIP-60 compliant URI to an OpenLineageDataset.
+
+    This function returns None if no URI normalizer is defined, no dataset converter is found or
+    some core Airflow changes are missing and ImportError is raised.
+    """
+    try:
+        from airflow.datasets import _get_normalized_scheme
+        from airflow.providers_manager import ProvidersManager
+
+        ol_converters = ProvidersManager().dataset_to_openlineage_converters
+        normalized_uri = dataset.normalized_uri
+    except (ImportError, AttributeError):
+        return None
+
+    if normalized_uri is None:
+        return None
+
+    if not (normalized_scheme := _get_normalized_scheme(normalized_uri)):
+        return None
+
+    if (airflow_to_ol_converter := ol_converters.get(normalized_scheme)) is None:
+        return None
+
+    return airflow_to_ol_converter(Dataset(uri=normalized_uri, extra=dataset.extra), lineage_context)

{apache_airflow_providers_openlineage-1.9.1.dist-info → apache_airflow_providers_openlineage-1.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-openlineage
-Version: 1.9.1
+Version: 1.10.0
 Summary: Provider package apache-airflow-providers-openlineage for Apache Airflow
 Keywords: airflow-provider,openlineage,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -28,8 +28,8 @@ Requires-Dist: openlineage-integration-common>=1.16.0
 Requires-Dist: openlineage-python>=1.16.0
 Requires-Dist: apache-airflow-providers-common-sql ; extra == "common.sql"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.10.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.10.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -80,7 +80,7 @@ Provides-Extra: common.sql
 
 Package ``apache-airflow-providers-openlineage``
 
-Release: ``1.9.1``
+Release: ``1.10.0``
 
 
 `OpenLineage <https://openlineage.io/>`__
@@ -93,7 +93,7 @@ This is a provider package for ``openlineage`` provider. All classes for this pr
 are in ``airflow.providers.openlineage`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.10.0/>`_.
 
 Installation
 ------------
@@ -137,4 +137,4 @@ Dependent package
 ============================================================================================================  ==============
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.10.0/changelog.html>`_.

apache_airflow_providers_openlineage-1.10.0.dist-info/RECORD

@@ -0,0 +1,29 @@
+airflow/providers/openlineage/LICENSE,sha256=FFb4jd2AXnOOf7XLP04pQW6jbdhG49TxlGY6fFpCV1Y,13609
+airflow/providers/openlineage/__init__.py,sha256=Zzf-0tKDYGhUI6kWk87QwKPT5SyjtGXripKvR-lSAak,1499
+airflow/providers/openlineage/conf.py,sha256=GcAY4poIXvFyhV3FywQFcgghl_XuMEIoVJOl3a_jKTU,5004
+airflow/providers/openlineage/get_provider_info.py,sha256=o02ioktk-8rLrEvtLx6pD7VfKLOn89iY8Ue9FPt5CDs,8359
+airflow/providers/openlineage/sqlparser.py,sha256=c7q3VVw41S87ZFozrkrEr2oZK79N12mC3KdDs9V1IuM,15581
+airflow/providers/openlineage/extractors/__init__.py,sha256=I0X4f6zUniclyD9zT0DFHRImpCpJVP4MkPJT3cd7X5I,1081
+airflow/providers/openlineage/extractors/base.py,sha256=W08Bg6FCSIzB9VLrm2fZ-zW0TLqfsIoY12hcD3qOzMw,6622
+airflow/providers/openlineage/extractors/bash.py,sha256=3aR0PXs8fzRLibRxXN1R8wMZnGzyCur7mjpy8e5GC4A,2583
+airflow/providers/openlineage/extractors/manager.py,sha256=uuSGLZgJHu0tu3RGzrc41uw6SahdgU3816T26FtnFEY,10539
+airflow/providers/openlineage/extractors/python.py,sha256=hVWOplMlBimrpPKPeW6vm75a8OmAYMU1oJzqMz8Jh90,3171
+airflow/providers/openlineage/facets/AirflowDagRunFacet.json,sha256=ie6c-J3-wGgk80WDTGWePz18o6DbW--TNM7BMF4WfcU,2251
+airflow/providers/openlineage/facets/AirflowJobFacet.json,sha256=rS9PuPWOi1Jc5B4a5qLxS_Az7Q9Eb3jVYQnN41iXDC0,1187
+airflow/providers/openlineage/facets/AirflowRunFacet.json,sha256=70mEaZShgSJp-2xr0bVvw3ljiGOPEaXD591fhuAQm_o,5953
+airflow/providers/openlineage/facets/AirflowStateRunFacet.json,sha256=xhHQEKD9Jopw-oqbkCCrrwFjfXnxvuJAritsmegKjuQ,937
+airflow/providers/openlineage/facets/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/openlineage/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/openlineage/plugins/adapter.py,sha256=4yKPGipMjn-E--RqGhTBzFCdBQjgmF0iBF1lSnw1p58,19430
+airflow/providers/openlineage/plugins/facets.py,sha256=1xeP1CCVSrVM_b_ehwlJ6JvW-0iZjYnbA73iWlif-4E,4500
+airflow/providers/openlineage/plugins/listener.py,sha256=MIXyHd-ITvNKVE7cij_v04MBtszJ-YQoRVM8s2DY3Fg,18435
+airflow/providers/openlineage/plugins/macros.py,sha256=hgFA3ZdQibyn4KXIOsKYBm4WRKDLA5q6Asscx5rvNfM,3076
+airflow/providers/openlineage/plugins/openlineage.py,sha256=rsRUW_zpXVAglzsgQRv5T9VWYY7CMQl0qRWm8-3oqDA,1678
+airflow/providers/openlineage/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/openlineage/utils/selective_enable.py,sha256=dFJ7wK7J_-BFwcOKp9tqFOSrASV3lmLv7HtRkEuMk3Q,3087
+airflow/providers/openlineage/utils/sql.py,sha256=bnuU9WvjVKcWVMN3cUp0jaHtU5_ZRM5I1OP1WhIdztg,9583
+airflow/providers/openlineage/utils/utils.py,sha256=O-jknMrWjzVvguCpBZERDfYny5wE6g6OepA3w9Bi0EQ,24557
+apache_airflow_providers_openlineage-1.10.0.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
+apache_airflow_providers_openlineage-1.10.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_openlineage-1.10.0.dist-info/METADATA,sha256=MhYQXkf5WMvmYd2aZi4QqnvpfRRbeD4JD2rSSYWbldQ,6374
+apache_airflow_providers_openlineage-1.10.0.dist-info/RECORD,,

apache_airflow_providers_openlineage-1.9.1.dist-info/RECORD

@@ -1,28 +0,0 @@
-airflow/providers/openlineage/LICENSE,sha256=FFb4jd2AXnOOf7XLP04pQW6jbdhG49TxlGY6fFpCV1Y,13609
-airflow/providers/openlineage/__init__.py,sha256=bDqRpjqfCE-UNXjjy8Ph6IFG8z5fdRrKJ9pOcX_nVuU,1498
-airflow/providers/openlineage/conf.py,sha256=ZHUZF3wjFGupfFa40F1DjZSkOQJCUtWFyQM3nvz6Bq4,5099
-airflow/providers/openlineage/get_provider_info.py,sha256=diZ2iRGa-47G5qbtylQ6gIXnN8rHXhTFmSK_mnJoaeE,7493
-airflow/providers/openlineage/sqlparser.py,sha256=qX2ID75PH11-youFvRdRHqZbYzLdbm95H3STvbBzGb8,15667
-airflow/providers/openlineage/extractors/__init__.py,sha256=I0X4f6zUniclyD9zT0DFHRImpCpJVP4MkPJT3cd7X5I,1081
-airflow/providers/openlineage/extractors/base.py,sha256=D-ll2cz_WWv1XYh2_UfzmpCl8S5iC2aptFTy2sWZeHY,5499
-airflow/providers/openlineage/extractors/bash.py,sha256=m4hLvDV4-zX4gp8apRuhpAR3Uakr8UOUxf-thTWmOxw,2563
-airflow/providers/openlineage/extractors/manager.py,sha256=9TyszMLAsgPS9NETWq7fPJjxcbTFk47x4kd1NRGCvsw,10315
-airflow/providers/openlineage/extractors/python.py,sha256=EQXCj2aHr2XXw0pNxeX-ii8UQFCoqkdf40ozqmA5d58,3151
-airflow/providers/openlineage/facets/AirflowJobFacet.json,sha256=rS9PuPWOi1Jc5B4a5qLxS_Az7Q9Eb3jVYQnN41iXDC0,1187
-airflow/providers/openlineage/facets/AirflowRunFacet.json,sha256=Mk85_KKwBOEcn_ptiQuoo9ucP-coAHfkAnilws2XMl4,5867
-airflow/providers/openlineage/facets/AirflowStateRunFacet.json,sha256=xhHQEKD9Jopw-oqbkCCrrwFjfXnxvuJAritsmegKjuQ,937
-airflow/providers/openlineage/facets/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/openlineage/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/openlineage/plugins/adapter.py,sha256=s0-eQw7cqzlUbznHdnLWihmdLKczHjuDj1TYYmJSK04,18720
-airflow/providers/openlineage/plugins/facets.py,sha256=kNxxnibLurcv8SV0qXPkwXxpKWhByo745vcqzWGNMIs,4450
-airflow/providers/openlineage/plugins/listener.py,sha256=q6r63SS3ZKdZKyTlEhw8c12onHh7kQCQ-_G1kA67btw,17952
-airflow/providers/openlineage/plugins/macros.py,sha256=hgFA3ZdQibyn4KXIOsKYBm4WRKDLA5q6Asscx5rvNfM,3076
-airflow/providers/openlineage/plugins/openlineage.py,sha256=rsRUW_zpXVAglzsgQRv5T9VWYY7CMQl0qRWm8-3oqDA,1678
-airflow/providers/openlineage/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/openlineage/utils/selective_enable.py,sha256=dFJ7wK7J_-BFwcOKp9tqFOSrASV3lmLv7HtRkEuMk3Q,3087
-airflow/providers/openlineage/utils/sql.py,sha256=5K3qFhevq-R4CDUZbq-XfFKPLWixIOtaypRem3jiHvw,9534
-airflow/providers/openlineage/utils/utils.py,sha256=cmPynb5QqZRV6Wzya-6cEcJUrCZjYm-4r4Z6NxwMzCc,19962
-apache_airflow_providers_openlineage-1.9.1.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
-apache_airflow_providers_openlineage-1.9.1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_openlineage-1.9.1.dist-info/METADATA,sha256=naW8pSs4fWp5BmV8wo5dbJa3buwt3uWWObAVfB0NMMg,6368
-apache_airflow_providers_openlineage-1.9.1.dist-info/RECORD,,