apache-airflow-providers-openlineage 1.8.0rc1__py3-none-any.whl → 1.9.0rc1__py3-none-any.whl

airflow/providers/openlineage/utils/utils.py

@@ -20,20 +20,27 @@ from __future__ import annotations
 import datetime
 import json
 import logging
-from contextlib import suppress
+import re
+from contextlib import redirect_stdout, suppress
 from functools import wraps
+from io import StringIO
 from typing import TYPE_CHECKING, Any, Iterable
 
 import attrs
 from deprecated import deprecated
 from openlineage.client.utils import RedactMixin
+from packaging.version import Version
 
+from airflow import __version__ as AIRFLOW_VERSION
 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 (
+    AirflowJobFacet,
     AirflowMappedTaskRunFacet,
     AirflowRunFacet,
+    AirflowStateRunFacet,
+    BaseFacet,
     UnknownOperatorAttributeRunFacet,
     UnknownOperatorInstance,
 )
@@ -41,6 +48,7 @@ from airflow.providers.openlineage.utils.selective_enable import (
     is_dag_lineage_enabled,
     is_task_lineage_enabled,
 )
+from airflow.serialization.serialized_objects import SerializedBaseOperator
 from airflow.utils.context import AirflowContextDeprecationWarning
 from airflow.utils.log.secrets_masker import Redactable, Redacted, SecretsMasker, should_hide_value_for_key
 from airflow.utils.module_loading import import_string
@@ -51,6 +59,7 @@ if TYPE_CHECKING:
 
 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")
 
 
 def try_import_from_string(string: str) -> Any:
@@ -78,7 +87,11 @@ def get_custom_facets(task_instance: TaskInstance | None = None) -> dict[str, An
 
 
 def get_fully_qualified_class_name(operator: BaseOperator | MappedOperator) -> str:
-    return operator.__class__.__module__ + "." + operator.__class__.__name__
+    if isinstance(operator, (MappedOperator, SerializedBaseOperator)):
+        # as in airflow.api_connexion.schemas.common_schema.ClassReferenceSchema
+        return operator._task_module + "." + operator._task_type  # type: ignore
+    op_class = get_operator_class(operator)
+    return op_class.__module__ + "." + op_class.__name__
 
 
 def is_operator_disabled(operator: BaseOperator | MappedOperator) -> bool:
@@ -135,6 +148,8 @@ class InfoJsonEncodable(dict):
     def _cast_basic_types(value):
         if isinstance(value, datetime.datetime):
             return value.isoformat()
+        if isinstance(value, datetime.timedelta):
+            return f"{value.total_seconds()} seconds"
         if isinstance(value, (set, list, tuple)):
             return str(list(value))
         return value
@@ -168,7 +183,7 @@ class InfoJsonEncodable(dict):
 class DagInfo(InfoJsonEncodable):
     """Defines encoding DAG object to JSON."""
 
-    includes = ["dag_id", "schedule_interval", "tags", "start_date"]
+    includes = ["dag_id", "description", "owner", "schedule_interval", "start_date", "tags"]
     casts = {"timetable": lambda dag: dag.timetable.serialize() if getattr(dag, "timetable", None) else None}
     renames = {"_dag_id": "dag_id"}
 
@@ -191,11 +206,11 @@ class DagRunInfo(InfoJsonEncodable):
 class TaskInstanceInfo(InfoJsonEncodable):
     """Defines encoding TaskInstance object to JSON."""
 
-    includes = ["duration", "try_number", "pool"]
+    includes = ["duration", "try_number", "pool", "queued_dttm"]
     casts = {
-        "map_index": lambda ti: ti.map_index
-        if hasattr(ti, "map_index") and getattr(ti, "map_index") != -1
-        else None
+        "map_index": lambda ti: (
+            ti.map_index if hasattr(ti, "map_index") and getattr(ti, "map_index") != -1 else None
+        )
     }
 
 
@@ -225,6 +240,7 @@ class TaskInfo(InfoJsonEncodable):
         "retries",
         "retry_exponential_backoff",
         "run_as_user",
+        "sla",
         "task_id",
         "trigger_rule",
         "upstream_task_ids",
@@ -234,9 +250,11 @@ class TaskInfo(InfoJsonEncodable):
     ]
     casts = {
         "operator_class": lambda task: task.task_type,
-        "task_group": lambda task: TaskGroupInfo(task.task_group)
-        if hasattr(task, "task_group") and getattr(task.task_group, "_group_id", None)
-        else None,
+        "task_group": lambda task: (
+            TaskGroupInfo(task.task_group)
+            if hasattr(task, "task_group") and getattr(task.task_group, "_group_id", None)
+            else None
+        ),
     }
 
 
@@ -262,20 +280,158 @@ def get_airflow_run_facet(
     task_instance: TaskInstance,
     task: BaseOperator,
     task_uuid: str,
-):
+) -> dict[str, BaseFacet]:
     return {
-        "airflow": attrs.asdict(
-            AirflowRunFacet(
-                dag=DagInfo(dag),
-                dagRun=DagRunInfo(dag_run),
-                taskInstance=TaskInstanceInfo(task_instance),
-                task=TaskInfo(task),
-                taskUuid=task_uuid,
-            )
+        "airflow": AirflowRunFacet(
+            dag=DagInfo(dag),
+            dagRun=DagRunInfo(dag_run),
+            taskInstance=TaskInstanceInfo(task_instance),
+            task=TaskInfo(task),
+            taskUuid=task_uuid,
+        )
+    }
+
+
+def get_airflow_job_facet(dag_run: DagRun) -> dict[str, BaseFacet]:
+    if not dag_run.dag:
+        return {}
+    return {
+        "airflow": AirflowJobFacet(
+            taskTree=_get_parsed_dag_tree(dag_run.dag),
+            taskGroups=_get_task_groups_details(dag_run.dag),
+            tasks=_get_tasks_details(dag_run.dag),
+        )
+    }
+
+
+def get_airflow_state_run_facet(dag_run: DagRun) -> dict[str, BaseFacet]:
+    return {
+        "airflowState": AirflowStateRunFacet(
+            dagRunState=dag_run.get_state(),
+            tasksState={ti.task_id: ti.state for ti in dag_run.get_task_instances()},
         )
     }
 
 
+def _safe_get_dag_tree_view(dag: DAG) -> list[str]:
+    # get_tree_view() has been added in Airflow 2.8.2
+    if hasattr(dag, "get_tree_view"):
+        return dag.get_tree_view().splitlines()
+
+    with redirect_stdout(StringIO()) as stdout:
+        dag.tree_view()
+        return stdout.getvalue().splitlines()
+
+
+def _get_parsed_dag_tree(dag: DAG) -> dict:
+    """
+    Get DAG's tasks hierarchy representation.
+
+    While the task dependencies are defined as following:
+    task >> [task_2, task_4] >> task_7
+    task_3 >> task_5
+    task_6  # has no dependencies, it's a root and a leaf
+
+    The result of this function will look like:
+    {
+        "task": {
+            "task_2": {
+                "task_7": {}
+            },
+            "task_4": {
+                "task_7": {}
+            }
+        },
+        "task_3": {
+            "task_5": {}
+        },
+        "task_6": {}
+    }
+    """
+    lines = _safe_get_dag_tree_view(dag)
+    task_dict: dict[str, dict] = {}
+    parent_map: dict[int, tuple[str, dict]] = {}
+
+    for line in lines:
+        stripped_line = line.strip()
+        if not stripped_line:
+            continue
+
+        # Determine the level by counting the leading spaces, assuming 4 spaces per level
+        # as defined in airflow.models.dag.DAG._generate_tree_view()
+        level = (len(line) - len(stripped_line)) // 4
+        # airflow.models.baseoperator.BaseOperator.__repr__ is used in DAG tree
+        # <Task({op_class}): {task_id}>
+        match = re.match(r"^<Task\((.+)\): (.*?)>$", stripped_line)
+        if not match:
+            return {}
+        current_task_id = match[2]
+
+        if level == 0:  # It's a root task
+            task_dict[current_task_id] = {}
+            parent_map[level] = (current_task_id, task_dict[current_task_id])
+        else:
+            # Find the immediate parent task
+            parent_task, parent_dict = parent_map[(level - 1)]
+            # Create new dict for the current task
+            parent_dict[current_task_id] = {}
+            # Update this task in the parent map
+            parent_map[level] = (current_task_id, parent_dict[current_task_id])
+
+    return task_dict
+
+
+def _get_tasks_details(dag: DAG) -> dict:
+    tasks = {
+        single_task.task_id: {
+            "operator": get_fully_qualified_class_name(single_task),
+            "task_group": single_task.task_group.group_id if single_task.task_group else None,
+            "emits_ol_events": _emits_ol_events(single_task),
+            "ui_color": single_task.ui_color,
+            "ui_fgcolor": single_task.ui_fgcolor,
+            "ui_label": single_task.label,
+            "is_setup": single_task.is_setup,
+            "is_teardown": single_task.is_teardown,
+        }
+        for single_task in dag.tasks
+    }
+
+    return tasks
+
+
+def _get_task_groups_details(dag: DAG) -> dict:
+    return {
+        tg_id: {
+            "parent_group": tg.parent_group.group_id,
+            "tooltip": tg.tooltip,
+            "ui_color": tg.ui_color,
+            "ui_fgcolor": tg.ui_fgcolor,
+            "ui_label": tg.label,
+        }
+        for tg_id, tg in dag.task_group_dict.items()
+    }
+
+
+def _emits_ol_events(task: BaseOperator | MappedOperator) -> bool:
+    config_selective_enabled = is_selective_lineage_enabled(task)
+    config_disabled_for_operators = is_operator_disabled(task)
+    # empty operators without callbacks/outlets are skipped for optimization by Airflow
+    # in airflow.models.taskinstance.TaskInstance._schedule_downstream_tasks
+    is_skipped_as_empty_operator = all(
+        (
+            task.inherits_from_empty_operator,
+            not task.on_execute_callback,
+            not task.on_success_callback,
+            not task.outlets,
+        )
+    )
+
+    emits_ol_events = all(
+        (config_selective_enabled, not config_disabled_for_operators, not is_skipped_as_empty_operator)
+    )
+    return emits_ol_events
+
+
 def get_unknown_source_attribute_run_facet(task: BaseOperator, name: str | None = None):
     if not name:
         name = get_operator_class(task).__name__
@@ -405,5 +561,7 @@ def normalize_sql(sql: str | Iterable[str]):
 
 
 def should_use_external_connection(hook) -> bool:
-    # TODO: Add checking overrides
-    return hook.__class__.__name__ not in ["SnowflakeHook", "SnowflakeSqlApiHook"]
+    # 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:
+        return hook.__class__.__name__ not in ["SnowflakeHook", "SnowflakeSqlApiHook", "RedshiftSQLHook"]
+    return True

{apache_airflow_providers_openlineage-1.8.0rc1.dist-info → apache_airflow_providers_openlineage-1.9.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-openlineage
-Version: 1.8.0rc1
+Version: 1.9.0rc1
 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>
@@ -24,12 +24,12 @@ Classifier: Topic :: System :: Monitoring
 Requires-Dist: apache-airflow-providers-common-sql>=1.6.0rc0
 Requires-Dist: apache-airflow>=2.7.0rc0
 Requires-Dist: attrs>=22.2
-Requires-Dist: openlineage-integration-common>=0.28.0
-Requires-Dist: openlineage-python>=0.28.0
+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.8.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.8.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.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.8.0.rc1``
+Release: ``1.9.0.rc1``
 
 
 `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.8.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.0/>`_.
 
 Installation
 ------------
@@ -113,8 +113,8 @@ PIP package                              Version required
 ``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``
+``openlineage-integration-common``       ``>=1.16.0``
+``openlineage-python``                   ``>=1.16.0``
 =======================================  ==================
 
 Cross provider package dependencies
@@ -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.8.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/1.9.0/changelog.html>`_.

apache_airflow_providers_openlineage-1.9.0rc1.dist-info/RECORD

@@ -0,0 +1,28 @@
+airflow/providers/openlineage/LICENSE,sha256=FFb4jd2AXnOOf7XLP04pQW6jbdhG49TxlGY6fFpCV1Y,13609
+airflow/providers/openlineage/__init__.py,sha256=wyNzY61gCpS2r6Wx-q-JiM8TbOKgTbDVngFZmcVmqOc,1498
+airflow/providers/openlineage/conf.py,sha256=ZHUZF3wjFGupfFa40F1DjZSkOQJCUtWFyQM3nvz6Bq4,5099
+airflow/providers/openlineage/get_provider_info.py,sha256=v-KK3oIOhEr-LE1WE8J6MaduGde709qiUmbkOklREn4,7472
+airflow/providers/openlineage/sqlparser.py,sha256=nMUA7Ak2HNfnGZ1PybWbNPENZfQIcJ1eE8AMh9UD-O8,15653
+airflow/providers/openlineage/extractors/__init__.py,sha256=I0X4f6zUniclyD9zT0DFHRImpCpJVP4MkPJT3cd7X5I,1081
+airflow/providers/openlineage/extractors/base.py,sha256=o6z8bXwNor1hwcUzezJ8LIPynR_BqXkP-qtwEgLtD2Q,5476
+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=nm8muyX-CwBaCNy_YDYbz-iFNL4msVZX6gFb4EaI-YE,18552
+airflow/providers/openlineage/plugins/facets.py,sha256=t60uCyp-FZKAeO7PSom5eAbn0cilN_Mfciip4FWG22M,4445
+airflow/providers/openlineage/plugins/listener.py,sha256=dk7omQItGSLcVw6bK56ScxLTnYQyQa1bXz2Fw402cjA,17956
+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=JVTmXdQknBL-9N0drFDkVMf1HCf8C6nbITVaP4-5ba4,3072
+airflow/providers/openlineage/utils/sql.py,sha256=yNBGMnWMECHtbpTh5nlaQh01O__LWkeQaSqrOVRWz6Y,9524
+airflow/providers/openlineage/utils/utils.py,sha256=QBg4B-isDrizb53cQFOflFx8LVKiabwE08I70I0AxOY,19485
+apache_airflow_providers_openlineage-1.9.0rc1.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
+apache_airflow_providers_openlineage-1.9.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_openlineage-1.9.0rc1.dist-info/METADATA,sha256=kamPdoPzxHZpb788jJ_RhWSQUgR1U76rKpRnraIm5Qg,6381
+apache_airflow_providers_openlineage-1.9.0rc1.dist-info/RECORD,,

apache_airflow_providers_openlineage-1.8.0rc1.dist-info/RECORD

@@ -1,24 +0,0 @@
-airflow/providers/openlineage/LICENSE,sha256=ywUBpKZc7Jb96rVt5I3IDbg7dIJAbUSHkuoDcF3jbH4,13569
-airflow/providers/openlineage/__init__.py,sha256=Y_3EiIS_TiqaVpc68HfekILHQRlsSGQLhs72joO7THg,1498
-airflow/providers/openlineage/conf.py,sha256=wozXzU5Do9S0mtjjGc5ruF556G2-ZT4GJa3YLT_-Phg,4693
-airflow/providers/openlineage/get_provider_info.py,sha256=ypUFlQXsC6s-cA7OdslwUaxhjaYIslbP_OILdb9CVNQ,7072
-airflow/providers/openlineage/sqlparser.py,sha256=-FGWWK0Xu6XkGSXcfn7PXsWIe0Y0fwe-3hivHg7emLA,15308
-airflow/providers/openlineage/extractors/__init__.py,sha256=I0X4f6zUniclyD9zT0DFHRImpCpJVP4MkPJT3cd7X5I,1081
-airflow/providers/openlineage/extractors/base.py,sha256=o6z8bXwNor1hwcUzezJ8LIPynR_BqXkP-qtwEgLtD2Q,5476
-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/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/openlineage/plugins/adapter.py,sha256=d_HczG_nXMwAHdGD-CO4baVfCc9ROcY0zQUWIZRF_Sw,16917
-airflow/providers/openlineage/plugins/facets.py,sha256=Z6dsz0rv-3VbRaGZJxW_T7Dak2k0ROGF0YrR_1awxZ0,2644
-airflow/providers/openlineage/plugins/listener.py,sha256=LrqvvGL4nzMNeha6PVbtylxp8gru3iKrkQltuQD2k8o,14092
-airflow/providers/openlineage/plugins/macros.py,sha256=QowPc9cc_unV-NLxBwm32OmWETA9pOucWguSeK92SSc,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=JVTmXdQknBL-9N0drFDkVMf1HCf8C6nbITVaP4-5ba4,3072
-airflow/providers/openlineage/utils/sql.py,sha256=7tEK0zVfIe7v3NI6oyv62x0KAS3sl8Ajfhqob8MdiX8,9366
-airflow/providers/openlineage/utils/utils.py,sha256=WFMdRsuArqqrgjsmFGb_ljIWV6ry-EGkTPZzy5aiG4Q,14036
-apache_airflow_providers_openlineage-1.8.0rc1.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
-apache_airflow_providers_openlineage-1.8.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_openlineage-1.8.0rc1.dist-info/METADATA,sha256=ljCFvcM-Hj_8eX0og8NVnFDxfQHfPPZYnzEa85ov8tI,6381
-apache_airflow_providers_openlineage-1.8.0rc1.dist-info/RECORD,,