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

airflow/providers/openlineage/plugins/facets.py

@@ -39,15 +39,56 @@ class AirflowMappedTaskRunFacet(BaseFacet):
 
     @classmethod
     def from_task_instance(cls, task_instance):
-        task = task_instance.task
-        from airflow.providers.openlineage.utils.utils import get_operator_class
+        from airflow.providers.openlineage.utils.utils import get_fully_qualified_class_name
 
         return cls(
             mapIndex=task_instance.map_index,
-            operatorClass=f"{get_operator_class(task).__module__}.{get_operator_class(task).__name__}",
+            operatorClass=get_fully_qualified_class_name(task_instance.task),
         )
 
 
+@define(slots=False)
+class AirflowJobFacet(BaseFacet):
+    """
+    Composite Airflow job facet.
+
+    This facet encapsulates all the necessary information to re-create full scope of an Airflow DAG logic,
+    enabling reconstruction, visualization, and analysis of DAGs in a comprehensive manner.
+    It includes detailed representations of the tasks, task groups, and their hierarchical relationships,
+    making it possible to draw a graph that visually represents the entire DAG structure (like in Airflow UI).
+    It also indicates whether a task should emit an OpenLineage (OL) event, enabling consumers to anticipate
+    the number of events and identify the tasks from which they can expect these events.
+
+    Attributes:
+        taskTree: A dictionary representing the hierarchical structure of tasks in the DAG.
+        taskGroups: A dictionary that contains information about task groups within the DAG.
+        tasks: A dictionary detailing individual tasks within the DAG.
+    """
+
+    taskTree: dict
+    taskGroups: dict
+    tasks: dict
+
+
+@define(slots=False)
+class AirflowStateRunFacet(BaseFacet):
+    """
+    Airflow facet providing state information.
+
+    This facet is designed to be sent at a completion event, offering state information about
+    the DAG run and each individual task. This information is crucial for understanding
+    the execution flow and comprehensive post-run analysis and debugging, including why certain tasks
+    did not emit events, which can occur due to the use of control flow operators like the BranchOperator.
+
+    Attributes:
+        dagRunState: This indicates the final status of the entire DAG run (e.g., "success", "failed").
+        tasksState: A dictionary mapping task IDs to their respective states. (e.g., "failed", "skipped").
+    """
+
+    dagRunState: str
+    tasksState: dict[str, str]
+
+
 @define(slots=False)
 class AirflowRunFacet(BaseFacet):
     """Composite Airflow run facet."""
@@ -61,7 +102,8 @@ class AirflowRunFacet(BaseFacet):
 
 @define(slots=False)
 class UnknownOperatorInstance(RedactMixin):
-    """Describes an unknown operator.
+    """
+    Describes an unknown operator.
 
     This specifies the (class) name of the operator and its properties.
     """

airflow/providers/openlineage/plugins/listener.py

@@ -17,18 +17,23 @@
 from __future__ import annotations
 
 import logging
+import os
 from concurrent.futures import ProcessPoolExecutor
 from datetime import datetime
 from typing import TYPE_CHECKING
 
+import psutil
 from openlineage.client.serde import Serde
+from packaging.version import Version
+from setproctitle import getproctitle, setproctitle
 
-from airflow import __version__ as airflow_version, settings
+from airflow import __version__ as AIRFLOW_VERSION, settings
 from airflow.listeners import hookimpl
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.extractors import ExtractorManager
 from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter, RunState
 from airflow.providers.openlineage.utils.utils import (
+    get_airflow_job_facet,
     get_airflow_run_facet,
     get_custom_facets,
     get_job_name,
@@ -36,6 +41,7 @@ from airflow.providers.openlineage.utils.utils import (
     is_selective_lineage_enabled,
     print_warning,
 )
+from airflow.settings import configure_orm
 from airflow.stats import Stats
 from airflow.utils.timeout import timeout
 
@@ -43,18 +49,29 @@ if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
     from airflow.models import DagRun, TaskInstance
+    from airflow.utils.state import TaskInstanceState
 
 _openlineage_listener: OpenLineageListener | None = None
+_IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
 
 
 def _get_try_number_success(val):
     # todo: remove when min airflow version >= 2.10.0
-    from packaging.version import parse
-
-    if parse(parse(airflow_version).base_version) < parse("2.10.0"):
-        return val.try_number - 1
-    else:
+    if _IS_AIRFLOW_2_10_OR_HIGHER:
         return val.try_number
+    return val.try_number - 1
+
+
+def _executor_initializer():
+    """
+    Initialize worker processes for the executor used for DagRun listener.
+
+    This function must be picklable, so it cannot be defined as an inner method or local function.
+
+    Reconfigures the ORM engine to prevent issues that arise when multiple processes interact with
+    the Airflow database.
+    """
+    settings.configure_orm()
 
 
 class OpenLineageListener:
@@ -69,10 +86,10 @@ class OpenLineageListener:
     @hookimpl
     def on_task_instance_running(
         self,
-        previous_state,
+        previous_state: TaskInstanceState,
         task_instance: TaskInstance,
         session: Session,  # This will always be QUEUED
-    ):
+    ) -> None:
         if not getattr(task_instance, "task", None) is not None:
             self.log.warning(
                 "No task set for TI object task_id: %s - dag_id: %s - run_id %s",
@@ -111,13 +128,16 @@ class OpenLineageListener:
             # we return here because Airflow 2.3 needs task from deferred state
             if task_instance.next_method is not None:
                 return
-            parent_run_id = self.adapter.build_dag_run_id(dag.dag_id, dagrun.run_id)
+            parent_run_id = self.adapter.build_dag_run_id(
+                dag_id=dag.dag_id,
+                execution_date=dagrun.execution_date,
+            )
 
             task_uuid = self.adapter.build_task_instance_run_id(
                 dag_id=dag.dag_id,
                 task_id=task.task_id,
-                execution_date=task_instance.execution_date,
                 try_number=task_instance.try_number,
+                execution_date=task_instance.execution_date,
             )
             event_type = RunState.RUNNING.value.lower()
             operator_name = task.task_type.lower()
@@ -130,7 +150,6 @@ class OpenLineageListener:
                 dagrun.data_interval_start.isoformat() if dagrun.data_interval_start else None
             )
             data_interval_end = dagrun.data_interval_end.isoformat() if dagrun.data_interval_end else None
-
             redacted_event = self.adapter.start_task(
                 run_id=task_uuid,
                 job_name=get_job_name(task),
@@ -153,10 +172,12 @@ class OpenLineageListener:
                 len(Serde.to_json(redacted_event).encode("utf-8")),
             )
 
-        on_running()
+        self._execute(on_running, "on_running", use_fork=True)
 
     @hookimpl
-    def on_task_instance_success(self, previous_state, task_instance: TaskInstance, session):
+    def on_task_instance_success(
+        self, previous_state: TaskInstanceState, task_instance: TaskInstance, session: Session
+    ) -> None:
         self.log.debug("OpenLineage listener got notification about task instance success")
 
         dagrun = task_instance.dag_run
@@ -184,13 +205,16 @@ class OpenLineageListener:
 
         @print_warning(self.log)
         def on_success():
-            parent_run_id = OpenLineageAdapter.build_dag_run_id(dag.dag_id, dagrun.run_id)
+            parent_run_id = OpenLineageAdapter.build_dag_run_id(
+                dag_id=dag.dag_id,
+                execution_date=dagrun.execution_date,
+            )
 
             task_uuid = OpenLineageAdapter.build_task_instance_run_id(
                 dag_id=dag.dag_id,
                 task_id=task.task_id,
-                execution_date=task_instance.execution_date,
                 try_number=_get_try_number_success(task_instance),
+                execution_date=task_instance.execution_date,
             )
             event_type = RunState.COMPLETE.value.lower()
             operator_name = task.task_type.lower()
@@ -215,10 +239,39 @@ class OpenLineageListener:
                 len(Serde.to_json(redacted_event).encode("utf-8")),
             )
 
-        on_success()
+        self._execute(on_success, "on_success", use_fork=True)
 
-    @hookimpl
-    def on_task_instance_failed(self, previous_state, task_instance: TaskInstance, session):
+    if _IS_AIRFLOW_2_10_OR_HIGHER:
+
+        @hookimpl
+        def on_task_instance_failed(
+            self,
+            previous_state: TaskInstanceState,
+            task_instance: TaskInstance,
+            error: None | str | BaseException,
+            session: Session,
+        ) -> None:
+            self._on_task_instance_failed(
+                previous_state=previous_state, task_instance=task_instance, error=error, session=session
+            )
+
+    else:
+
+        @hookimpl
+        def on_task_instance_failed(
+            self, previous_state: TaskInstanceState, task_instance: TaskInstance, session: Session
+        ) -> None:
+            self._on_task_instance_failed(
+                previous_state=previous_state, task_instance=task_instance, error=None, session=session
+            )
+
+    def _on_task_instance_failed(
+        self,
+        previous_state: TaskInstanceState,
+        task_instance: TaskInstance,
+        session: Session,
+        error: None | str | BaseException = None,
+    ) -> None:
         self.log.debug("OpenLineage listener got notification about task instance failure")
 
         dagrun = task_instance.dag_run
@@ -246,13 +299,16 @@ class OpenLineageListener:
 
         @print_warning(self.log)
         def on_failure():
-            parent_run_id = OpenLineageAdapter.build_dag_run_id(dag.dag_id, dagrun.run_id)
+            parent_run_id = OpenLineageAdapter.build_dag_run_id(
+                dag_id=dag.dag_id,
+                execution_date=dagrun.execution_date,
+            )
 
             task_uuid = OpenLineageAdapter.build_task_instance_run_id(
                 dag_id=dag.dag_id,
                 task_id=task.task_id,
-                execution_date=task_instance.execution_date,
                 try_number=task_instance.try_number,
+                execution_date=task_instance.execution_date,
             )
             event_type = RunState.FAIL.value.lower()
             operator_name = task.task_type.lower()
@@ -271,40 +327,76 @@ class OpenLineageListener:
                 parent_run_id=parent_run_id,
                 end_time=end_date.isoformat(),
                 task=task_metadata,
+                error=error,
             )
             Stats.gauge(
                 f"ol.event.size.{event_type}.{operator_name}",
                 len(Serde.to_json(redacted_event).encode("utf-8")),
             )
 
-        on_failure()
+        self._execute(on_failure, "on_failure", use_fork=True)
+
+    def _execute(self, callable, callable_name: str, use_fork: bool = False):
+        if use_fork:
+            self._fork_execute(callable, callable_name)
+        else:
+            callable()
+
+    def _terminate_with_wait(self, process: psutil.Process):
+        process.terminate()
+        try:
+            # Waiting for max 3 seconds to make sure process can clean up before being killed.
+            process.wait(timeout=3)
+        except psutil.TimeoutExpired:
+            # If it's not dead by then, then force kill.
+            process.kill()
+
+    def _fork_execute(self, callable, callable_name: str):
+        self.log.debug("Will fork to execute OpenLineage process.")
+        pid = os.fork()
+        if pid:
+            process = psutil.Process(pid)
+            try:
+                self.log.debug("Waiting for process %s", pid)
+                process.wait(conf.execution_timeout())
+            except psutil.TimeoutExpired:
+                self.log.warning(
+                    "OpenLineage process %s expired. This should not affect process execution.", pid
+                )
+                self._terminate_with_wait(process)
+            except BaseException:
+                # Kill the process directly.
+                self._terminate_with_wait(process)
+            self.log.debug("Process with pid %s finished - parent", pid)
+        else:
+            setproctitle(getproctitle() + " - OpenLineage - " + callable_name)
+            configure_orm(disable_connection_pool=True)
+            self.log.debug("Executing OpenLineage process - %s - pid %s", callable_name, os.getpid())
+            callable()
+            self.log.debug("Process with current pid finishes after %s", callable_name)
+            os._exit(0)
 
     @property
-    def executor(self):
-        def initializer():
-            # Re-configure the ORM engine as there are issues with multiple processes
-            # if process calls Airflow DB.
-            settings.configure_orm()
-
+    def executor(self) -> ProcessPoolExecutor:
         if not self._executor:
             self._executor = ProcessPoolExecutor(
                 max_workers=conf.dag_state_change_process_pool_size(),
-                initializer=initializer,
+                initializer=_executor_initializer,
             )
         return self._executor
 
     @hookimpl
-    def on_starting(self, component):
+    def on_starting(self, component) -> None:
         self.log.debug("on_starting: %s", component.__class__.__name__)
 
     @hookimpl
-    def before_stopping(self, component):
+    def before_stopping(self, component) -> None:
         self.log.debug("before_stopping: %s", component.__class__.__name__)
         with timeout(30):
             self.executor.shutdown(wait=True)
 
     @hookimpl
-    def on_dag_run_running(self, dag_run: DagRun, msg: str):
+    def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None:
         if dag_run.dag and not is_selective_lineage_enabled(dag_run.dag):
             self.log.debug(
                 "Skipping OpenLineage event emission for DAG `%s` "
@@ -326,10 +418,13 @@ class OpenLineageListener:
             msg=msg,
             nominal_start_time=data_interval_start,
             nominal_end_time=data_interval_end,
+            # AirflowJobFacet should be created outside ProcessPoolExecutor that pickles objects,
+            # as it causes lack of some TaskGroup attributes and crashes event emission.
+            job_facets={**get_airflow_job_facet(dag_run=dag_run)},
         )
 
     @hookimpl
-    def on_dag_run_success(self, dag_run: DagRun, msg: str):
+    def on_dag_run_success(self, dag_run: DagRun, msg: str) -> None:
         if dag_run.dag and not is_selective_lineage_enabled(dag_run.dag):
             self.log.debug(
                 "Skipping OpenLineage event emission for DAG `%s` "
@@ -346,7 +441,7 @@ class OpenLineageListener:
         self.executor.submit(self.adapter.dag_success, dag_run=dag_run, msg=msg)
 
     @hookimpl
-    def on_dag_run_failed(self, dag_run: DagRun, msg: str):
+    def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None:
         if dag_run.dag and not is_selective_lineage_enabled(dag_run.dag):
             self.log.debug(
                 "Skipping OpenLineage event emission for DAG `%s` "

airflow/providers/openlineage/plugins/macros.py

@@ -61,8 +61,8 @@ def lineage_run_id(task_instance: TaskInstance):
     return OpenLineageAdapter.build_task_instance_run_id(
         dag_id=task_instance.dag_id,
         task_id=task_instance.task_id,
-        execution_date=task_instance.execution_date,
         try_number=task_instance.try_number,
+        execution_date=task_instance.execution_date,
     )
 
 

airflow/providers/openlineage/sqlparser.py

@@ -39,6 +39,7 @@ from airflow.providers.openlineage.utils.sql import (
     get_table_schemas,
 )
 from airflow.typing_compat import TypedDict
+from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Engine
@@ -116,19 +117,27 @@ def from_table_meta(
     return Dataset(namespace=namespace, name=name if not is_uppercase else name.upper())
 
 
-class SQLParser:
-    """Interface for openlineage-sql.
+class SQLParser(LoggingMixin):
+    """
+    Interface for openlineage-sql.
 
     :param dialect: dialect specific to the database
     :param default_schema: schema applied to each table with no schema parsed
     """
 
     def __init__(self, dialect: str | None = None, default_schema: str | None = None) -> None:
+        super().__init__()
         self.dialect = dialect
         self.default_schema = default_schema
 
     def parse(self, sql: list[str] | str) -> SqlMeta | None:
         """Parse a single or a list of SQL statements."""
+        self.log.debug(
+            "OpenLineage calling SQL parser with SQL %s dialect %s schema %s",
+            sql,
+            self.dialect,
+            self.default_schema,
+        )
         return parse(sql=sql, dialect=self.dialect, default_schema=self.default_schema)
 
     def parse_table_schemas(
@@ -151,6 +160,7 @@ class SQLParser:
             "database": database or database_info.database,
             "use_flat_cross_db_query": database_info.use_flat_cross_db_query,
         }
+        self.log.info("PRE getting schemas for input and output tables")
         return get_table_schemas(
             hook,
             namespace,
@@ -235,7 +245,8 @@ class SQLParser:
         sqlalchemy_engine: Engine | None = None,
         use_connection: bool = True,
     ) -> OperatorLineage:
-        """Parse SQL statement(s) and generate OpenLineage metadata.
+        """
+        Parse SQL statement(s) and generate OpenLineage metadata.
 
         Generated OpenLineage metadata contains:
 
@@ -335,9 +346,8 @@ class SQLParser:
             return split_statement(sql)
         return [obj for stmt in sql for obj in cls.split_sql_string(stmt) if obj != ""]
 
-    @classmethod
     def create_information_schema_query(
-        cls,
+        self,
         tables: list[DbTableMeta],
         normalize_name: Callable[[str], str],
         is_cross_db: bool,
@@ -349,7 +359,7 @@ class SQLParser:
         sqlalchemy_engine: Engine | None = None,
     ) -> str:
         """Create SELECT statement to query information schema table."""
-        tables_hierarchy = cls._get_tables_hierarchy(
+        tables_hierarchy = self._get_tables_hierarchy(
             tables,
             normalize_name=normalize_name,
             database=database,

airflow/providers/openlineage/utils/selective_enable.py

@@ -32,7 +32,8 @@ log = logging.getLogger(__name__)
 
 
 def enable_lineage(obj: T) -> T:
-    """Set selective enable OpenLineage parameter to True.
+    """
+    Set selective enable OpenLineage parameter to True.
 
     The method also propagates param to tasks if the object is DAG.
     """
@@ -48,7 +49,8 @@ def enable_lineage(obj: T) -> T:
 
 
 def disable_lineage(obj: T) -> T:
-    """Set selective enable OpenLineage parameter to False.
+    """
+    Set selective enable OpenLineage parameter to False.
 
     The method also propagates param to tasks if the object is DAG.
     """
@@ -73,7 +75,8 @@ def is_task_lineage_enabled(task: Operator) -> bool:
 
 
 def is_dag_lineage_enabled(dag: DAG) -> bool:
-    """Check if DAG is selectively enabled to emit OpenLineage events.
+    """
+    Check if DAG is selectively enabled to emit OpenLineage events.
 
     The method also checks if selective enable parameter is set to True
     or if any of the tasks in DAG is selectively enabled.

airflow/providers/openlineage/utils/sql.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import logging
 from collections import defaultdict
 from contextlib import closing
 from enum import IntEnum
@@ -33,6 +34,9 @@ if TYPE_CHECKING:
     from airflow.hooks.base import BaseHook
 
 
+log = logging.getLogger(__name__)
+
+
 class ColumnIndex(IntEnum):
     """Enumerates the indices of columns in information schema view."""
 
@@ -81,7 +85,8 @@ def get_table_schemas(
     in_query: str | None,
     out_query: str | None,
 ) -> tuple[list[Dataset], list[Dataset]]:
-    """Query database for table schemas.
+    """
+    Query database for table schemas.
 
     Uses provided hook. Responsibility to provide queries for this function is on particular extractors.
     If query for input or output table isn't provided, the query is skipped.
@@ -90,6 +95,7 @@ def get_table_schemas(
     if not in_query and not out_query:
         return [], []
 
+    log.debug("Starting to query database for table schemas")
     with closing(hook.get_conn()) as conn, closing(conn.cursor()) as cursor:
         if in_query:
             cursor.execute(in_query)
@@ -101,11 +107,13 @@ def get_table_schemas(
             out_datasets = [x.to_dataset(namespace, database, schema) for x in parse_query_result(cursor)]
         else:
             out_datasets = []
+    log.debug("Got table schema query result from database.")
     return in_datasets, out_datasets
 
 
 def parse_query_result(cursor) -> list[TableSchema]:
-    """Fetch results from DB-API 2.0 cursor and creates list of table schemas.
+    """
+    Fetch results from DB-API 2.0 cursor and creates list of table schemas.
 
     For each row it creates :class:`TableSchema`.
     """
@@ -149,7 +157,7 @@ def create_information_schema_query(
     sqlalchemy_engine: Engine | None = None,
 ) -> str:
     """Create query for getting table schemas from information schema."""
-    metadata = MetaData(sqlalchemy_engine)
+    metadata = MetaData()
     select_statements = []
     # Don't iterate over tables hierarchy, just pass it to query single information schema table
     if use_flat_cross_db_query: