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

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.warning("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,7 +117,7 @@ def from_table_meta(
     return Dataset(namespace=namespace, name=name if not is_uppercase else name.upper())
 
 
-class SQLParser:
+class SQLParser(LoggingMixin):
     """Interface for openlineage-sql.
 
     :param dialect: dialect specific to the database
@@ -124,11 +125,18 @@ class SQLParser:
     """
 
     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 +159,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,
@@ -335,9 +344,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 +357,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/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."""
 
@@ -90,6 +94,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,6 +106,7 @@ 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
 
 
@@ -149,7 +155,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: