apache-airflow-providers-openlineage 2.1.3__py3-none-any.whl → 2.2.0__py3-none-any.whl

airflow/providers/openlineage/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "2.1.3"
+__version__ = "2.2.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.9.0"

airflow/providers/openlineage/extractors/manager.py

@@ -298,12 +298,11 @@ class ExtractorManager(LoggingMixin):
 
         if isinstance(obj, Dataset):
             return obj
-        elif isinstance(obj, Table):
+        if isinstance(obj, Table):
             return ExtractorManager.convert_to_ol_dataset_from_table(obj)
-        elif isinstance(obj, File):
+        if isinstance(obj, File):
             return ExtractorManager.convert_to_ol_dataset_from_object_storage_uri(obj.url)
-        else:
-            return None
+        return None
 
     def validate_task_metadata(self, task_metadata) -> OperatorLineage | None:
         try:

airflow/providers/openlineage/plugins/adapter.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import os
 import traceback
 from contextlib import ExitStack
 from typing import TYPE_CHECKING
@@ -36,6 +37,7 @@ from openlineage.client.facet_v2 import (
 )
 from openlineage.client.uuid import generate_static_uuid
 
+from airflow.configuration import conf as airflow_conf
 from airflow.providers.openlineage import __version__ as OPENLINEAGE_PROVIDER_VERSION, conf
 from airflow.providers.openlineage.utils.utils import (
     OpenLineageRedactor,
@@ -81,6 +83,11 @@ class OpenLineageAdapter(LoggingMixin):
 
     def get_or_create_openlineage_client(self) -> OpenLineageClient:
         if not self._client:
+            # If not already set explicitly - propagate airflow logging level to OpenLineage client
+            airflow_core_log_level = airflow_conf.get("logging", "logging_level", fallback="INFO")
+            if not os.getenv("OPENLINEAGE_CLIENT_LOGGING") and airflow_core_log_level != "INFO":
+                os.environ["OPENLINEAGE_CLIENT_LOGGING"] = airflow_core_log_level
+
             config = self.get_openlineage_config()
             if config:
                 self.log.debug(
@@ -102,8 +109,7 @@ class OpenLineageAdapter(LoggingMixin):
         if openlineage_config_path:
             config = self._read_yaml_config(openlineage_config_path)
             return config
-        else:
-            self.log.debug("OpenLineage config_path configuration not found.")
+        self.log.debug("OpenLineage config_path configuration not found.")
 
         # Second, try to get transport config
         transport_config = conf.transport()

airflow/providers/openlineage/plugins/listener.py

@@ -28,9 +28,9 @@ from setproctitle import getproctitle, setproctitle
 
 from airflow import settings
 from airflow.listeners import hookimpl
-from airflow.models import DagRun
+from airflow.models import DagRun, TaskInstance
 from airflow.providers.openlineage import conf
-from airflow.providers.openlineage.extractors import ExtractorManager
+from airflow.providers.openlineage.extractors import ExtractorManager, OperatorLineage
 from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter, RunState
 from airflow.providers.openlineage.utils.utils import (
     AIRFLOW_V_2_10_PLUS,
@@ -53,7 +53,6 @@ from airflow.utils.state import TaskInstanceState
 from airflow.utils.timeout import timeout
 
 if TYPE_CHECKING:
-    from airflow.models import TaskInstance
     from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
     from airflow.settings import Session
 
@@ -235,10 +234,18 @@ class OpenLineageListener:
 
         @hookimpl
         def on_task_instance_success(
-            self, previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance
+            self, previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance | TaskInstance
         ) -> None:
             self.log.debug("OpenLineage listener got notification about task instance success")
 
+            if isinstance(task_instance, TaskInstance):
+                self._on_task_instance_manual_state_change(
+                    ti=task_instance,
+                    dagrun=task_instance.dag_run,
+                    ti_state=TaskInstanceState.SUCCESS,
+                )
+                return
+
             context = task_instance.get_template_context()
             task = context["task"]
             if TYPE_CHECKING:
@@ -338,10 +345,20 @@ class OpenLineageListener:
         def on_task_instance_failed(
             self,
             previous_state: TaskInstanceState,
-            task_instance: TaskInstance,
+            task_instance: RuntimeTaskInstance | TaskInstance,
             error: None | str | BaseException,
         ) -> None:
             self.log.debug("OpenLineage listener got notification about task instance failure")
+
+            if isinstance(task_instance, TaskInstance):
+                self._on_task_instance_manual_state_change(
+                    ti=task_instance,
+                    dagrun=task_instance.dag_run,
+                    ti_state=TaskInstanceState.FAILED,
+                    error=error,
+                )
+                return
+
             context = task_instance.get_template_context()
             task = context["task"]
             if TYPE_CHECKING:
@@ -457,6 +474,60 @@ class OpenLineageListener:
 
         self._execute(on_failure, "on_failure", use_fork=True)
 
+    def _on_task_instance_manual_state_change(
+        self,
+        ti: TaskInstance,
+        dagrun: DagRun,
+        ti_state: TaskInstanceState,
+        error: None | str | BaseException = None,
+    ) -> None:
+        self.log.debug("`_on_task_instance_manual_state_change` was called with state: `%s`.", ti_state)
+        end_date = timezone.utcnow()
+
+        @print_warning(self.log)
+        def on_state_change():
+            date = dagrun.logical_date or dagrun.run_after
+            parent_run_id = self.adapter.build_dag_run_id(
+                dag_id=dagrun.dag_id,
+                logical_date=date,
+                clear_number=dagrun.clear_number,
+            )
+
+            task_uuid = self.adapter.build_task_instance_run_id(
+                dag_id=dagrun.dag_id,
+                task_id=ti.task_id,
+                try_number=ti.try_number,
+                logical_date=date,
+                map_index=ti.map_index,
+            )
+
+            adapter_kwargs = {
+                "run_id": task_uuid,
+                "job_name": get_job_name(ti),
+                "parent_job_name": dagrun.dag_id,
+                "parent_run_id": parent_run_id,
+                "end_time": end_date.isoformat(),
+                "task": OperatorLineage(),
+                "run_facets": get_airflow_debug_facet(),
+            }
+
+            if ti_state == TaskInstanceState.FAILED:
+                event_type = RunState.FAIL.value.lower()
+                redacted_event = self.adapter.fail_task(**adapter_kwargs, error=error)
+            elif ti_state == TaskInstanceState.SUCCESS:
+                event_type = RunState.COMPLETE.value.lower()
+                redacted_event = self.adapter.complete_task(**adapter_kwargs)
+            else:
+                raise ValueError(f"Unsupported ti_state: `{ti_state}`.")
+
+            operator_name = ti.operator.lower()
+            Stats.gauge(
+                f"ol.event.size.{event_type}.{operator_name}",
+                len(Serde.to_json(redacted_event).encode("utf-8")),
+            )
+
+        self._execute(on_state_change, "on_state_change", use_fork=True)
+
     def _execute(self, callable, callable_name: str, use_fork: bool = False):
         if use_fork:
             self._fork_execute(callable, callable_name)

airflow/providers/openlineage/utils/spark.py

@@ -111,7 +111,7 @@ def _get_transport_information_as_spark_properties() -> dict:
             props = _format_transport(props, http_transport, name)
         return props
 
-    elif transport.kind == "http":
+    if transport.kind == "http":
         return _format_transport({}, _get_transport_information(transport), None)
 
     log.info(

airflow/providers/openlineage/utils/utils.py

@@ -54,7 +54,9 @@ from airflow.providers.openlineage.version_compat import AIRFLOW_V_2_10_PLUS, AI
 from airflow.sensors.base import BaseSensorOperator
 from airflow.serialization.serialized_objects import SerializedBaseOperator
 from airflow.utils.module_loading import import_string
-from airflow.utils.session import NEW_SESSION, provide_session
+
+if not AIRFLOW_V_3_0_PLUS:
+    from airflow.utils.session import NEW_SESSION, provide_session
 
 try:
     from airflow.sdk import BaseOperator as SdkBaseOperator
@@ -195,40 +197,41 @@ def is_selective_lineage_enabled(obj: DAG | BaseOperator | MappedOperator | SdkB
         return True
     if isinstance(obj, DAG):
         return is_dag_lineage_enabled(obj)
-    elif isinstance(obj, (BaseOperator, MappedOperator, SdkBaseOperator)):
+    if isinstance(obj, (BaseOperator, MappedOperator, SdkBaseOperator)):
         return is_task_lineage_enabled(obj)
-    else:
-        raise TypeError("is_selective_lineage_enabled can only be used on DAG or Operator objects")
+    raise TypeError("is_selective_lineage_enabled can only be used on DAG or Operator objects")
 
 
-@provide_session
-def is_ti_rescheduled_already(ti: TaskInstance, session=NEW_SESSION):
-    from sqlalchemy import exists
+if not AIRFLOW_V_3_0_PLUS:
 
-    if not isinstance(ti.task, BaseSensorOperator):
-        return False
+    @provide_session
+    def is_ti_rescheduled_already(ti: TaskInstance, session=NEW_SESSION):
+        from sqlalchemy import exists
 
-    if not ti.task.reschedule:
-        return False
-    if AIRFLOW_V_3_0_PLUS:
+        if not isinstance(ti.task, BaseSensorOperator):
+            return False
+
+        if not ti.task.reschedule:
+            return False
+        if AIRFLOW_V_3_0_PLUS:
+            return (
+                session.query(
+                    exists().where(TaskReschedule.ti_id == ti.id, TaskReschedule.try_number == ti.try_number)
+                ).scalar()
+                is True
+            )
         return (
             session.query(
-                exists().where(TaskReschedule.ti_id == ti.id, TaskReschedule.try_number == ti.try_number)
+                exists().where(
+                    TaskReschedule.dag_id == ti.dag_id,
+                    TaskReschedule.task_id == ti.task_id,
+                    TaskReschedule.run_id == ti.run_id,
+                    TaskReschedule.map_index == ti.map_index,
+                    TaskReschedule.try_number == ti.try_number,
+                )
             ).scalar()
             is True
         )
-    return (
-        session.query(
-            exists().where(
-                TaskReschedule.dag_id == ti.dag_id,
-                TaskReschedule.task_id == ti.task_id,
-                TaskReschedule.run_id == ti.run_id,
-                TaskReschedule.map_index == ti.map_index,
-                TaskReschedule.try_number == ti.try_number,
-            )
-        ).scalar()
-        is True
-    )
 
 
 class InfoJsonEncodable(dict):
@@ -322,6 +325,7 @@ class DagInfo(InfoJsonEncodable):
         "description",
         "fileloc",
         "owner",
+        "owner_links",
         "schedule_interval",  # For Airflow 2.
         "timetable_summary",  # For Airflow 3.
         "start_date",
@@ -713,7 +717,7 @@ class OpenLineageRedactor(SecretsMasker):
                                 ),
                             )
                     return item
-                elif is_json_serializable(item) and hasattr(item, "__dict__"):
+                if is_json_serializable(item) and hasattr(item, "__dict__"):
                     for dict_key, subval in item.__dict__.items():
                         if type(subval).__name__ == "Proxy":
                             return "<<non-redactable: Proxy>>"
@@ -729,8 +733,7 @@ class OpenLineageRedactor(SecretsMasker):
                                 ),
                             )
                     return item
-                else:
-                    return super()._redact(item, name, depth, max_depth)
+                return super()._redact(item, name, depth, max_depth)
         except Exception as exc:
             log.warning("Unable to redact %r. Error was: %s: %s", item, type(exc).__name__, exc)
         return item

{apache_airflow_providers_openlineage-2.1.3.dist-info → apache_airflow_providers_openlineage-2.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-openlineage
-Version: 2.1.3
+Version: 2.2.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>
@@ -24,11 +24,11 @@ Requires-Dist: apache-airflow>=2.9.0
 Requires-Dist: apache-airflow-providers-common-sql>=1.20.0
 Requires-Dist: apache-airflow-providers-common-compat>=1.4.0
 Requires-Dist: attrs>=22.2
-Requires-Dist: openlineage-integration-common>=1.24.2
-Requires-Dist: openlineage-python>=1.24.2
+Requires-Dist: openlineage-integration-common>=1.31.0
+Requires-Dist: openlineage-python>=1.31.0
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.3/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.3
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.2.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.2.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -59,7 +59,7 @@ Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 Package ``apache-airflow-providers-openlineage``
 
-Release: ``2.1.3``
+Release: ``2.2.0``
 
 
 `OpenLineage <https://openlineage.io/>`__
@@ -72,7 +72,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/2.1.3/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.2.0/>`_.
 
 Installation
 ------------
@@ -93,8 +93,8 @@ PIP package                                 Version required
 ``apache-airflow-providers-common-sql``     ``>=1.20.0``
 ``apache-airflow-providers-common-compat``  ``>=1.4.0``
 ``attrs``                                   ``>=22.2``
-``openlineage-integration-common``          ``>=1.24.2``
-``openlineage-python``                      ``>=1.24.2``
+``openlineage-integration-common``          ``>=1.31.0``
+``openlineage-python``                      ``>=1.31.0``
 ==========================================  ==================
 
 Cross provider package dependencies
@@ -118,5 +118,5 @@ Dependent package
 ==================================================================================================================  =================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.3/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.2.0/changelog.html>`_.
 

{apache_airflow_providers_openlineage-2.1.3.dist-info → apache_airflow_providers_openlineage-2.2.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
 airflow/providers/openlineage/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-airflow/providers/openlineage/__init__.py,sha256=R6VkEEY6inOTMg3aCbbydjy-qIpvgikyWt7OGFyLts4,1498
+airflow/providers/openlineage/__init__.py,sha256=HaGAryTFkC-tcSplPNX4LM4SyZ52yNLtPbPcaLgk2dE,1498
 airflow/providers/openlineage/conf.py,sha256=aYdLU7iHBdGIU8ZAC5iUiIDgXP9gvP9r_z5hTAbXPOU,5535
 airflow/providers/openlineage/get_provider_info.py,sha256=108mAg-tdcTBmPYwJ9M3wcmhc-i1bAs5OB2MLnSBA-4,9256
 airflow/providers/openlineage/sqlparser.py,sha256=N38XhkU-lxwxnYevQpq63JOBi4rzp0q56JjxO3H24W8,20340
@@ -7,7 +7,7 @@ airflow/providers/openlineage/version_compat.py,sha256=aHg90_DtgoSnQvILFICexMyNl
 airflow/providers/openlineage/extractors/__init__.py,sha256=I0X4f6zUniclyD9zT0DFHRImpCpJVP4MkPJT3cd7X5I,1081
 airflow/providers/openlineage/extractors/base.py,sha256=ZXRlvMSak8kUfur-BxrgAxeylMQFG-iT-LusQguIFLc,6342
 airflow/providers/openlineage/extractors/bash.py,sha256=3aR0PXs8fzRLibRxXN1R8wMZnGzyCur7mjpy8e5GC4A,2583
-airflow/providers/openlineage/extractors/manager.py,sha256=g3WJRBR2-XZHTG7qAR4UEviwtymvDArhlPMVD3c4q_g,12862
+airflow/providers/openlineage/extractors/manager.py,sha256=kVo5OHkpUoYIiT2RvwUt128jC_Q4EosdLC0sP1YfIH0,12840
 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/AirflowDebugRunFacet.json,sha256=_zA5gFqGje5MOH1SmdMeA5ViOHvW_pV4oijEAvkuBbY,768
@@ -16,17 +16,17 @@ airflow/providers/openlineage/facets/AirflowRunFacet.json,sha256=70mEaZShgSJp-2x
 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=wCGJ3rMDpiyFlPaonUQwGzs0hyNqU-4__e0Z9nAbcyI,20620
+airflow/providers/openlineage/plugins/adapter.py,sha256=wSqdbxNhDWzGd4J6SmNFI_Hj7uMN9I7VZMPTTu0u0-w,21046
 airflow/providers/openlineage/plugins/facets.py,sha256=VvyMYR6ONkC95q5FdNmohv0scbA1Ej_B5cQ97as5GvA,4161
-airflow/providers/openlineage/plugins/listener.py,sha256=KlBKT9VkdOrZxvQHsLZWWq_g4jPhaa2GdVxmHy_EVhM,26083
+airflow/providers/openlineage/plugins/listener.py,sha256=J3TB6o8TFYvLlK2um4tNWGPSjGoZHdXkDhbpzIQzBD0,28845
 airflow/providers/openlineage/plugins/macros.py,sha256=qrHLjE95Uq8H-W9CIkQe5Y9Pu1O-GErhpDV2olGaGQM,3730
 airflow/providers/openlineage/plugins/openlineage.py,sha256=HD3mYNPfXd-buZydEpuAY-naVBXhausU2LYUNhL48QA,1906
 airflow/providers/openlineage/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/openlineage/utils/selective_enable.py,sha256=YyrUQ7Djv5o46XdH83N_G8AXAZ9C_aKPa534pbNVp08,3441
-airflow/providers/openlineage/utils/spark.py,sha256=-2XfUaV0WISK6vHSBmB9E78xkuPjO3fM1tDQCZG7j9I,7303
+airflow/providers/openlineage/utils/spark.py,sha256=sOfEkw6OgFbVG8w-PlIlWqHdq8ti_ja7b3bJnEesWiM,7301
 airflow/providers/openlineage/utils/sql.py,sha256=vkKrrdENEMVG8gtzV6yuTXMa2Z9fBAEXmxDVIDaVncI,9571
-airflow/providers/openlineage/utils/utils.py,sha256=_TP49gO4GeyXqqqRgYY-7bgNaM1cczLTv-3lQFGLMfM,29329
-apache_airflow_providers_openlineage-2.1.3.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
-apache_airflow_providers_openlineage-2.1.3.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_openlineage-2.1.3.dist-info/METADATA,sha256=4KKyiTnqidPb-eRmkd_CrEQwi0r6iW98e51OF26xzfM,5685
-apache_airflow_providers_openlineage-2.1.3.dist-info/RECORD,,
+airflow/providers/openlineage/utils/utils.py,sha256=qDAwGpYaTKQokkT40iQgPOjCuL2ZIPLufNuCAalH7-0,29472
+apache_airflow_providers_openlineage-2.2.0.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
+apache_airflow_providers_openlineage-2.2.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_openlineage-2.2.0.dist-info/METADATA,sha256=6Nf9R8Oo0Z3Jw9qYU43KzFzOB-t7pPvTadv7V1h-n_0,5685
+apache_airflow_providers_openlineage-2.2.0.dist-info/RECORD,,