apache-airflow-providers-openlineage 2.0.0rc2__py3-none-any.whl → 2.1.0rc1__py3-none-any.whl

airflow/providers/openlineage/plugins/listener.py

@@ -19,10 +19,10 @@ 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 setproctitle import getproctitle, setproctitle
 
 from airflow import settings
@@ -33,6 +33,7 @@ from airflow.providers.openlineage.extractors import ExtractorManager
 from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter, RunState
 from airflow.providers.openlineage.utils.utils import (
     AIRFLOW_V_2_10_PLUS,
+    AIRFLOW_V_3_0_PLUS,
     get_airflow_dag_run_facet,
     get_airflow_debug_facet,
     get_airflow_job_facet,
@@ -42,7 +43,6 @@ from airflow.providers.openlineage.utils.utils import (
     get_user_provided_run_facets,
     is_operator_disabled,
     is_selective_lineage_enabled,
-    is_ti_rescheduled_already,
     print_warning,
 )
 from airflow.settings import configure_orm
@@ -50,11 +50,12 @@ from airflow.stats import Stats
 from airflow.utils import timezone
 from airflow.utils.state import TaskInstanceState
 from airflow.utils.timeout import timeout
+from openlineage.client.serde import Serde
 
 if TYPE_CHECKING:
-    from sqlalchemy.orm import Session
-
     from airflow.models import TaskInstance
+    from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+    from airflow.settings import Session
 
 _openlineage_listener: OpenLineageListener | None = None
 
@@ -87,28 +88,58 @@ class OpenLineageListener:
         self.extractor_manager = ExtractorManager()
         self.adapter = OpenLineageAdapter()
 
-    @hookimpl
-    def on_task_instance_running(
-        self,
-        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",
-                task_instance.task_id,
-                task_instance.dag_id,
-                task_instance.run_id,
-            )
-            return
+    if AIRFLOW_V_3_0_PLUS:
 
-        self.log.debug("OpenLineage listener got notification about task instance start")
-        dagrun = task_instance.dag_run
-        task = task_instance.task
-        if TYPE_CHECKING:
-            assert task
-        dag = task.dag
+        @hookimpl
+        def on_task_instance_running(
+            self,
+            previous_state: TaskInstanceState,
+            task_instance: RuntimeTaskInstance,
+        ):
+            self.log.debug("OpenLineage listener got notification about task instance start")
+            context = task_instance.get_template_context()
+
+            task = context["task"]
+            if TYPE_CHECKING:
+                assert task
+            dagrun = context["dag_run"]
+            dag = context["dag"]
+            start_date = context["start_date"]
+            self._on_task_instance_running(task_instance, dag, dagrun, task, start_date)
+    else:
+
+        @hookimpl
+        def on_task_instance_running(
+            self,
+            previous_state: TaskInstanceState,
+            task_instance: TaskInstance,
+            session: Session,  # type: ignore[valid-type]
+        ) -> None:
+            from airflow.providers.openlineage.utils.utils import is_ti_rescheduled_already
+
+            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",
+                    task_instance.task_id,
+                    task_instance.dag_id,
+                    task_instance.run_id,
+                )
+                return
+
+            self.log.debug("OpenLineage listener got notification about task instance start")
+            task = task_instance.task
+            if TYPE_CHECKING:
+                assert task
+            start_date = task_instance.start_date if task_instance.start_date else timezone.utcnow()
+
+            if is_ti_rescheduled_already(task_instance):
+                self.log.debug("Skipping this instance of rescheduled task - START event was emitted already")
+                return
+            self._on_task_instance_running(task_instance, task.dag, task_instance.dag_run, task, start_date)
+
+    def _on_task_instance_running(
+        self, task_instance: RuntimeTaskInstance | TaskInstance, dag, dagrun, task, start_date: datetime
+    ):
         if is_operator_disabled(task):
             self.log.debug(
                 "Skipping OpenLineage event emission for operator `%s` "
@@ -127,35 +158,41 @@ class OpenLineageListener:
             return
 
         # Needs to be calculated outside of inner method so that it gets cached for usage in fork processes
+        data_interval_start = dagrun.data_interval_start
+        if isinstance(data_interval_start, datetime):
+            data_interval_start = data_interval_start.isoformat()
+        data_interval_end = dagrun.data_interval_end
+        if isinstance(data_interval_end, datetime):
+            data_interval_end = data_interval_end.isoformat()
+
+        clear_number = 0
+        if hasattr(dagrun, "clear_number"):
+            clear_number = dagrun.clear_number
+
         debug_facet = get_airflow_debug_facet()
 
         @print_warning(self.log)
         def on_running():
-            # that's a workaround to detect task running from deferred state
-            # we return here because Airflow 2.3 needs task from deferred state
-            if task_instance.next_method is not None:
-                return
-
-            if is_ti_rescheduled_already(task_instance):
+            context = task_instance.get_template_context()
+            if hasattr(context, "task_reschedule_count") and context["task_reschedule_count"] > 0:
                 self.log.debug("Skipping this instance of rescheduled task - START event was emitted already")
                 return
 
+            date = dagrun.logical_date
+            if AIRFLOW_V_3_0_PLUS and date is None:
+                date = dagrun.run_after
+
             parent_run_id = self.adapter.build_dag_run_id(
                 dag_id=dag.dag_id,
-                logical_date=dagrun.logical_date,
-                clear_number=dagrun.clear_number,
+                logical_date=date,
+                clear_number=clear_number,
             )
 
-            if hasattr(task_instance, "logical_date"):
-                logical_date = task_instance.logical_date
-            else:
-                logical_date = task_instance.execution_date
-
             task_uuid = self.adapter.build_task_instance_run_id(
                 dag_id=dag.dag_id,
                 task_id=task.task_id,
                 try_number=task_instance.try_number,
-                logical_date=logical_date,
+                logical_date=date,
                 map_index=task_instance.map_index,
             )
             event_type = RunState.RUNNING.value.lower()
@@ -164,11 +201,6 @@ class OpenLineageListener:
             with Stats.timer(f"ol.extract.{event_type}.{operator_name}"):
                 task_metadata = self.extractor_manager.extract_metadata(dagrun, task)
 
-            start_date = task_instance.start_date if task_instance.start_date else timezone.utcnow()
-            data_interval_start = (
-                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),
@@ -195,17 +227,39 @@ class OpenLineageListener:
 
         self._execute(on_running, "on_running", use_fork=True)
 
-    @hookimpl
-    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")
+    if AIRFLOW_V_3_0_PLUS:
+
+        @hookimpl
+        def on_task_instance_success(
+            self, previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance
+        ) -> None:
+            self.log.debug("OpenLineage listener got notification about task instance success")
 
-        dagrun = task_instance.dag_run
-        task = task_instance.task
-        if TYPE_CHECKING:
-            assert task
-        dag = task.dag
+            context = task_instance.get_template_context()
+            task = context["task"]
+            if TYPE_CHECKING:
+                assert task
+            dagrun = context["dag_run"]
+            dag = context["dag"]
+            self._on_task_instance_success(task_instance, dag, dagrun, task)
+
+    else:
+
+        @hookimpl
+        def on_task_instance_success(
+            self,
+            previous_state: TaskInstanceState,
+            task_instance: TaskInstance,
+            session: Session,  # type: ignore[valid-type]
+        ) -> None:
+            self.log.debug("OpenLineage listener got notification about task instance success")
+            task = task_instance.task
+            if TYPE_CHECKING:
+                assert task
+            self._on_task_instance_success(task_instance, task.dag, task_instance.dag_run, task)
+
+    def _on_task_instance_success(self, task_instance: RuntimeTaskInstance, dag, dagrun, task):
+        end_date = timezone.utcnow()
 
         if is_operator_disabled(task):
             self.log.debug(
@@ -226,21 +280,21 @@ class OpenLineageListener:
 
         @print_warning(self.log)
         def on_success():
+            date = dagrun.logical_date
+            if AIRFLOW_V_3_0_PLUS and date is None:
+                date = dagrun.run_after
+
             parent_run_id = self.adapter.build_dag_run_id(
                 dag_id=dag.dag_id,
-                logical_date=dagrun.logical_date,
+                logical_date=date,
                 clear_number=dagrun.clear_number,
             )
 
-            if hasattr(task_instance, "logical_date"):
-                logical_date = task_instance.logical_date
-            else:
-                logical_date = task_instance.execution_date
             task_uuid = self.adapter.build_task_instance_run_id(
                 dag_id=dag.dag_id,
                 task_id=task.task_id,
                 try_number=_get_try_number_success(task_instance),
-                logical_date=logical_date,
+                logical_date=date,
                 map_index=task_instance.map_index,
             )
             event_type = RunState.COMPLETE.value.lower()
@@ -251,8 +305,6 @@ class OpenLineageListener:
                     dagrun, task, complete=True, task_instance=task_instance
                 )
 
-            end_date = task_instance.end_date if task_instance.end_date else timezone.utcnow()
-
             redacted_event = self.adapter.complete_task(
                 run_id=task_uuid,
                 job_name=get_job_name(task),
@@ -273,7 +325,7 @@ class OpenLineageListener:
 
         self._execute(on_success, "on_success", use_fork=True)
 
-    if AIRFLOW_V_2_10_PLUS:
+    if AIRFLOW_V_3_0_PLUS:
 
         @hookimpl
         def on_task_instance_failed(
@@ -281,36 +333,54 @@ class OpenLineageListener:
             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
-            )
+            self.log.debug("OpenLineage listener got notification about task instance failure")
+            context = task_instance.get_template_context()
+            task = context["task"]
+            if TYPE_CHECKING:
+                assert task
+            dagrun = context["dag_run"]
+            dag = context["dag"]
+            self._on_task_instance_failed(task_instance, dag, dagrun, task, error)
 
+    elif AIRFLOW_V_2_10_PLUS:
+
+        @hookimpl
+        def on_task_instance_failed(
+            self,
+            previous_state: TaskInstanceState,
+            task_instance: TaskInstance,
+            error: None | str | BaseException,
+            session: Session,  # type: ignore[valid-type]
+        ) -> None:
+            self.log.debug("OpenLineage listener got notification about task instance failure")
+            task = task_instance.task
+            if TYPE_CHECKING:
+                assert task
+            self._on_task_instance_failed(task_instance, task.dag, task_instance.dag_run, task, error)
     else:
 
         @hookimpl
         def on_task_instance_failed(
-            self, previous_state: TaskInstanceState, task_instance: TaskInstance, session: Session
+            self,
+            previous_state: TaskInstanceState,
+            task_instance: TaskInstance,
+            session: Session,  # type: ignore[valid-type]
         ) -> None:
-            self._on_task_instance_failed(
-                previous_state=previous_state, task_instance=task_instance, error=None, session=session
-            )
+            task = task_instance.task
+            if TYPE_CHECKING:
+                assert task
+            self._on_task_instance_failed(task_instance, task.dag, task_instance.dag_run, task)
 
     def _on_task_instance_failed(
         self,
-        previous_state: TaskInstanceState,
-        task_instance: TaskInstance,
-        session: Session,
+        task_instance: TaskInstance | RuntimeTaskInstance,
+        dag,
+        dagrun,
+        task,
         error: None | str | BaseException = None,
     ) -> None:
-        self.log.debug("OpenLineage listener got notification about task instance failure")
-
-        dagrun = task_instance.dag_run
-        task = task_instance.task
-        if TYPE_CHECKING:
-            assert task
-        dag = task.dag
+        end_date = timezone.utcnow()
 
         if is_operator_disabled(task):
             self.log.debug(
@@ -331,22 +401,21 @@ class OpenLineageListener:
 
         @print_warning(self.log)
         def on_failure():
+            date = dagrun.logical_date
+            if AIRFLOW_V_3_0_PLUS and date is None:
+                date = dagrun.run_after
+
             parent_run_id = self.adapter.build_dag_run_id(
                 dag_id=dag.dag_id,
-                logical_date=dagrun.logical_date,
+                logical_date=date,
                 clear_number=dagrun.clear_number,
             )
 
-            if hasattr(task_instance, "logical_date"):
-                logical_date = task_instance.logical_date
-            else:
-                logical_date = task_instance.execution_date
-
             task_uuid = self.adapter.build_task_instance_run_id(
                 dag_id=dag.dag_id,
                 task_id=task.task_id,
                 try_number=task_instance.try_number,
-                logical_date=logical_date,
+                logical_date=date,
                 map_index=task_instance.map_index,
             )
             event_type = RunState.FAIL.value.lower()
@@ -357,8 +426,6 @@ class OpenLineageListener:
                     dagrun, task, complete=True, task_instance=task_instance
                 )
 
-            end_date = task_instance.end_date if task_instance.end_date else timezone.utcnow()
-
             redacted_event = self.adapter.fail_task(
                 run_id=task_uuid,
                 job_name=get_job_name(task),
@@ -462,10 +529,14 @@ class OpenLineageListener:
 
             run_facets = {**get_airflow_dag_run_facet(dag_run)}
 
+            date = dag_run.logical_date
+            if AIRFLOW_V_3_0_PLUS and date is None:
+                date = dag_run.run_after
+
             self.submit_callable(
                 self.adapter.dag_started,
                 dag_id=dag_run.dag_id,
-                logical_date=dag_run.logical_date,
+                logical_date=date,
                 start_date=dag_run.start_date,
                 nominal_start_time=data_interval_start,
                 nominal_end_time=data_interval_end,
@@ -500,15 +571,21 @@ class OpenLineageListener:
                 task_ids = DagRun._get_partial_task_ids(dag_run.dag)
             else:
                 task_ids = dag_run.dag.task_ids if dag_run.dag and dag_run.dag.partial else None
+
+            date = dag_run.logical_date
+            if AIRFLOW_V_3_0_PLUS and date is None:
+                date = dag_run.run_after
+
             self.submit_callable(
                 self.adapter.dag_success,
                 dag_id=dag_run.dag_id,
                 run_id=dag_run.run_id,
                 end_date=dag_run.end_date,
-                logical_date=dag_run.logical_date,
+                logical_date=date,
                 clear_number=dag_run.clear_number,
                 task_ids=task_ids,
                 dag_run_state=dag_run.get_state(),
+                run_facets={**get_airflow_dag_run_facet(dag_run)},
             )
         except BaseException as e:
             self.log.warning("OpenLineage received exception in method on_dag_run_success", exc_info=e)
@@ -533,16 +610,22 @@ class OpenLineageListener:
                 task_ids = DagRun._get_partial_task_ids(dag_run.dag)
             else:
                 task_ids = dag_run.dag.task_ids if dag_run.dag and dag_run.dag.partial else None
+
+            date = dag_run.logical_date
+            if AIRFLOW_V_3_0_PLUS and date is None:
+                date = dag_run.run_after
+
             self.submit_callable(
                 self.adapter.dag_failed,
                 dag_id=dag_run.dag_id,
                 run_id=dag_run.run_id,
                 end_date=dag_run.end_date,
-                logical_date=dag_run.logical_date,
+                logical_date=date,
                 clear_number=dag_run.clear_number,
                 dag_run_state=dag_run.get_state(),
                 task_ids=task_ids,
                 msg=msg,
+                run_facets={**get_airflow_dag_run_facet(dag_run)},
             )
         except BaseException as e:
             self.log.warning("OpenLineage received exception in method on_dag_run_failed", exc_info=e)

airflow/providers/openlineage/sqlparser.py

@@ -16,13 +16,11 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Callable
+import logging
+from typing import TYPE_CHECKING, Callable, TypedDict
 
 import sqlparse
 from attrs import define
-from openlineage.client.event_v2 import Dataset
-from openlineage.client.facet_v2 import column_lineage_dataset, extraction_error_run, sql_job
-from openlineage.common.sql import DbTableMeta, SqlMeta, parse
 
 from airflow.providers.openlineage.extractors.base import OperatorLineage
 from airflow.providers.openlineage.utils.sql import (
@@ -30,14 +28,20 @@ from airflow.providers.openlineage.utils.sql import (
     create_information_schema_query,
     get_table_schemas,
 )
-from airflow.typing_compat import TypedDict
+from airflow.providers.openlineage.utils.utils import should_use_external_connection
 from airflow.utils.log.logging_mixin import LoggingMixin
+from openlineage.client.event_v2 import Dataset
+from openlineage.client.facet_v2 import column_lineage_dataset, extraction_error_run, sql_job
+from openlineage.common.sql import DbTableMeta, SqlMeta, parse
 
 if TYPE_CHECKING:
-    from openlineage.client.facet_v2 import JobFacet, RunFacet
     from sqlalchemy.engine import Engine
 
     from airflow.hooks.base import BaseHook
+    from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from openlineage.client.facet_v2 import JobFacet, RunFacet
+
+log = logging.getLogger(__name__)
 
 DEFAULT_NAMESPACE = "default"
 DEFAULT_INFORMATION_SCHEMA_COLUMNS = [
@@ -78,10 +82,67 @@ class DatabaseInfo:
     :param database: Takes precedence over parsed database name.
     :param information_schema_columns: List of columns names from information schema table.
     :param information_schema_table_name: Information schema table name.
-    :param use_flat_cross_db_query: Specifies if single information schema table should be used
-        for cross-database queries (e.g. for Redshift).
-    :param is_information_schema_cross_db: Specifies if information schema contains
-        cross-database data.
+    :param use_flat_cross_db_query: Specifies whether a single, "global" information schema table should
+        be used for cross-database queries (e.g., in Redshift), or if multiple, per-database "local"
+        information schema tables should be queried individually.
+
+        If True, assumes a single, universal information schema table is available
+        (for example, in Redshift, the `SVV_REDSHIFT_COLUMNS` view)
+        [https://docs.aws.amazon.com/redshift/latest/dg/r_SVV_REDSHIFT_COLUMNS.html].
+        In this mode, we query only `information_schema_table_name` directly.
+        Depending on the `is_information_schema_cross_db` argument, you can also filter
+        by database name in the WHERE clause.
+
+        If False, treats each database as having its own local information schema table containing
+        metadata for that database only. As a result, one query per database may be generated
+        and then combined (often via `UNION ALL`).
+        This approach is necessary for dialects that do not maintain a single global view of
+        all metadata or that require per-database queries.
+        Depending on the `is_information_schema_cross_db` argument, queries can
+        include or omit database information in both identifiers and filters.
+
+        See `is_information_schema_cross_db` which also affects how final queries are constructed.
+    :param is_information_schema_cross_db: Specifies whether database information should be tracked
+        and included in queries that retrieve schema information from the information_schema_table.
+        In short, this determines whether queries are capable of spanning multiple databases.
+
+        If True, database identifiers are included wherever applicable, allowing retrieval of
+        metadata from more than one database. For instance, in Snowflake or MS SQL
+        (where each database is treated as a top-level namespace), you might have a query like:
+
+        ```
+        SELECT ...
+        FROM db1.information_schema.columns WHERE ...
+        UNION ALL
+        SELECT ...
+        FROM db2.information_schema.columns WHERE ...
+        ```
+
+        In Redshift, setting this to True together with `use_flat_cross_db_query=True` allows
+        adding database filters to the query, for example:
+
+        ```
+        SELECT ...
+        FROM SVV_REDSHIFT_COLUMNS
+        WHERE
+        SVV_REDSHIFT_COLUMNS.database == db1  # This is skipped when False
+        AND SVV_REDSHIFT_COLUMNS.schema == schema1
+        AND SVV_REDSHIFT_COLUMNS.table IN (table1, table2)
+        OR ...
+        ```
+
+        However, certain databases (e.g., PostgreSQL) do not permit true cross-database queries.
+        In such dialects, enabling cross-database support may lead to errors or be unnecessary.
+        Always consult your dialect's documentation or test sample queries to confirm if
+        cross-database querying is supported.
+
+        If False, database qualifiers are ignored, effectively restricting queries to a single
+        database (or making the database-level qualifier optional). This is typically
+        safer for databases that do not support cross-database operations or only provide a
+        two-level namespace (schema + table) instead of a three-level one (database + schema + table).
+        For example, some MySQL or PostgreSQL contexts might not need or permit cross-database queries at all.
+
+        See `use_flat_cross_db_query` which also affects how final queries are constructed.
     :param is_uppercase_names: Specifies if database accepts only uppercase names (e.g. Snowflake).
     :param normalize_name_method: Method to normalize database, schema and table names.
         Defaults to `name.lower()`.
@@ -397,3 +458,43 @@ class SQLParser(LoggingMixin):
             tables = schemas.setdefault(normalize_name(table.schema) if table.schema else None, [])
             tables.append(table.name)
         return hierarchy
+
+
+def get_openlineage_facets_with_sql(
+    hook: DbApiHook, sql: str | list[str], conn_id: str, database: str | None
+) -> OperatorLineage | None:
+    connection = hook.get_connection(conn_id)
+    try:
+        database_info = hook.get_openlineage_database_info(connection)
+    except AttributeError:
+        database_info = None
+
+    if database_info is None:
+        log.debug("%s has no database info provided", hook)
+        return None
+
+    try:
+        sql_parser = SQLParser(
+            dialect=hook.get_openlineage_database_dialect(connection),
+            default_schema=hook.get_openlineage_default_schema(),
+        )
+    except AttributeError:
+        log.debug("%s failed to get database dialect", hook)
+        return None
+
+    try:
+        sqlalchemy_engine = hook.get_sqlalchemy_engine()
+    except Exception as e:
+        log.debug("Failed to get sql alchemy engine: %s", e)
+        sqlalchemy_engine = None
+
+    operator_lineage = sql_parser.generate_openlineage_metadata_from_sql(
+        sql=sql,
+        hook=hook,
+        database_info=database_info,
+        database=database,
+        sqlalchemy_engine=sqlalchemy_engine,
+        use_connection=should_use_external_connection(hook),
+    )
+
+    return operator_lineage

airflow/providers/openlineage/utils/selective_enable.py

@@ -18,16 +18,29 @@
 from __future__ import annotations
 
 import logging
-from typing import TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
-from airflow.models import DAG, Operator, Param
+from airflow.models import Param
 from airflow.models.xcom_arg import XComArg
 
+if TYPE_CHECKING:
+    from airflow.sdk import DAG
+    from airflow.sdk.definitions._internal.abstractoperator import Operator
+else:
+    try:
+        from airflow.sdk import DAG
+    except ImportError:
+        from airflow.models import DAG
+
 ENABLE_OL_PARAM_NAME = "_selective_enable_ol"
 ENABLE_OL_PARAM = Param(True, const=True)
 DISABLE_OL_PARAM = Param(False, const=False)
 T = TypeVar("T", bound="DAG | Operator")
 
+if TYPE_CHECKING:
+    from airflow.sdk.definitions.baseoperator import BaseOperator as SdkBaseOperator
+
+
 log = logging.getLogger(__name__)
 
 
@@ -65,7 +78,7 @@ def disable_lineage(obj: T) -> T:
     return obj
 
 
-def is_task_lineage_enabled(task: Operator) -> bool:
+def is_task_lineage_enabled(task: Operator | SdkBaseOperator) -> bool:
     """Check if selective enable OpenLineage parameter is set to True on task level."""
     if task.params.get(ENABLE_OL_PARAM_NAME) is False:
         log.debug(