apache-airflow-providers-openlineage 2.1.1rc1__tar.gz → 2.1.2rc1__tar.gz

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-openlineage
-Version: 2.1.1rc1
+Version: 2.1.2rc1
 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>
@@ -27,11 +27,11 @@ Requires-Dist: attrs>=22.2
 Requires-Dist: openlineage-integration-common>=1.24.2
 Requires-Dist: openlineage-python>=1.24.2
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2
+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
-Project-URL: Twitter, https://x.com/ApacheAirflow
 Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 
@@ -59,7 +59,7 @@ Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 Package ``apache-airflow-providers-openlineage``
 
-Release: ``2.1.1``
+Release: ``2.1.2``
 
 
 `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.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2/>`_.
 
 Installation
 ------------
@@ -101,7 +101,7 @@ Cross provider package dependencies
 -----------------------------------
 
 Those are dependencies that might be needed in order to use all the features of the package.
-You need to install the specified provider packages in order to use them.
+You need to install the specified providers in order to use them.
 
 You can install such cross-provider dependencies when installing from PyPI. For example:
 
@@ -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.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2/changelog.html>`_.
 

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-openlineage``
 
-Release: ``2.1.1``
+Release: ``2.1.2``
 
 
 `OpenLineage <https://openlineage.io/>`__
@@ -36,7 +36,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.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2/>`_.
 
 Installation
 ------------
@@ -65,7 +65,7 @@ Cross provider package dependencies
 -----------------------------------
 
 Those are dependencies that might be needed in order to use all the features of the package.
-You need to install the specified provider packages in order to use them.
+You need to install the specified providers in order to use them.
 
 You can install such cross-provider dependencies when installing from PyPI. For example:
 
@@ -82,4 +82,4 @@ 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.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2/changelog.html>`_.

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/pyproject.toml

@@ -20,12 +20,12 @@
 # IF YOU WANT TO MODIFY THIS FILE EXCEPT DEPENDENCIES, YOU SHOULD MODIFY THE TEMPLATE
 # `pyproject_TEMPLATE.toml.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
 [build-system]
-requires = ["flit_core==3.11.0"]
+requires = ["flit_core==3.12.0"]
 build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-openlineage"
-version = "2.1.1.rc1"
+version = "2.1.2.rc1"
 description = "Provider package apache-airflow-providers-openlineage for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -87,12 +87,12 @@ apache-airflow-providers-fab = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.1"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.1/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-openlineage/2.1.2/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"
-"Twitter" = "https://x.com/ApacheAirflow"
+"Mastodon" = "https://fosstodon.org/@airflow"
 "YouTube" = "https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/"
 
 [project.entry-points."apache_airflow_provider"]

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/__init__.py

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

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/extractors/base.py

@@ -29,14 +29,16 @@ with warnings.catch_warnings():
     from openlineage.client.facet import BaseFacet as BaseFacet_V1
 from openlineage.client.facet_v2 import JobFacet, RunFacet
 
-from airflow.providers.openlineage.utils.utils import AIRFLOW_V_2_10_PLUS
 from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.state import TaskInstanceState
 
 # this is not to break static checks compatibility with v1 OpenLineage facet classes
 DatasetSubclass = TypeVar("DatasetSubclass", bound=OLDataset)
 BaseFacetSubclass = TypeVar("BaseFacetSubclass", bound=Union[BaseFacet_V1, RunFacet, JobFacet])
 
+OL_METHOD_NAME_START = "get_openlineage_facets_on_start"
+OL_METHOD_NAME_COMPLETE = "get_openlineage_facets_on_complete"
+OL_METHOD_NAME_FAIL = "get_openlineage_facets_on_failure"
+
 
 @define
 class OperatorLineage(Generic[DatasetSubclass, BaseFacetSubclass]):
@@ -81,6 +83,9 @@ class BaseExtractor(ABC, LoggingMixin):
     def extract_on_complete(self, task_instance) -> OperatorLineage | None:
         return self.extract()
 
+    def extract_on_failure(self, task_instance) -> OperatorLineage | None:
+        return self.extract()
+
 
 class DefaultExtractor(BaseExtractor):
     """Extractor that uses `get_openlineage_facets_on_start/complete/failure` methods."""
@@ -96,46 +101,41 @@ class DefaultExtractor(BaseExtractor):
         return []
 
     def _execute_extraction(self) -> OperatorLineage | None:
-        # OpenLineage methods are optional - if there's no method, return None
-        try:
-            self.log.debug(
-                "Trying to execute `get_openlineage_facets_on_start` for %s.", self.operator.task_type
-            )
-            return self._get_openlineage_facets(self.operator.get_openlineage_facets_on_start)  # type: ignore
-        except ImportError:
-            self.log.error(
-                "OpenLineage provider method failed to import OpenLineage integration. "
-                "This should not happen. Please report this bug to developers."
-            )
-            return None
-        except AttributeError:
+        method = getattr(self.operator, OL_METHOD_NAME_START, None)
+        if callable(method):
             self.log.debug(
-                "Operator %s does not have the get_openlineage_facets_on_start method.",
-                self.operator.task_type,
+                "Trying to execute '%s' method of '%s'.", OL_METHOD_NAME_START, self.operator.task_type
             )
-            return OperatorLineage()
+            return self._get_openlineage_facets(method)
+        self.log.debug(
+            "Operator '%s' does not have '%s' method.", self.operator.task_type, OL_METHOD_NAME_START
+        )
+        return OperatorLineage()
 
     def extract_on_complete(self, task_instance) -> OperatorLineage | None:
-        failed_states = [TaskInstanceState.FAILED, TaskInstanceState.UP_FOR_RETRY]
-        if not AIRFLOW_V_2_10_PLUS:  # todo: remove when min airflow version >= 2.10.0
-            # Before fix (#41053) implemented in Airflow 2.10 TaskInstance's state was still RUNNING when
-            # being passed to listener's on_failure method. Since `extract_on_complete()` is only called
-            # after task completion, RUNNING state means that we are dealing with FAILED task in < 2.10
-            failed_states = [TaskInstanceState.RUNNING]
-
-        if task_instance.state in failed_states:
-            on_failed = getattr(self.operator, "get_openlineage_facets_on_failure", None)
-            if on_failed and callable(on_failed):
-                self.log.debug(
-                    "Executing `get_openlineage_facets_on_failure` for %s.", self.operator.task_type
-                )
-                return self._get_openlineage_facets(on_failed, task_instance)
-        on_complete = getattr(self.operator, "get_openlineage_facets_on_complete", None)
-        if on_complete and callable(on_complete):
-            self.log.debug("Executing `get_openlineage_facets_on_complete` for %s.", self.operator.task_type)
-            return self._get_openlineage_facets(on_complete, task_instance)
+        method = getattr(self.operator, OL_METHOD_NAME_COMPLETE, None)
+        if callable(method):
+            self.log.debug(
+                "Trying to execute '%s' method of '%s'.", OL_METHOD_NAME_COMPLETE, self.operator.task_type
+            )
+            return self._get_openlineage_facets(method, task_instance)
+        self.log.debug(
+            "Operator '%s' does not have '%s' method.", self.operator.task_type, OL_METHOD_NAME_COMPLETE
+        )
         return self.extract()
 
+    def extract_on_failure(self, task_instance) -> OperatorLineage | None:
+        method = getattr(self.operator, OL_METHOD_NAME_FAIL, None)
+        if callable(method):
+            self.log.debug(
+                "Trying to execute '%s' method of '%s'.", OL_METHOD_NAME_FAIL, self.operator.task_type
+            )
+            return self._get_openlineage_facets(method, task_instance)
+        self.log.debug(
+            "Operator '%s' does not have '%s' method.", self.operator.task_type, OL_METHOD_NAME_FAIL
+        )
+        return self.extract_on_complete(task_instance)
+
     def _get_openlineage_facets(self, get_facets_method, *args) -> OperatorLineage | None:
         try:
             facets: OperatorLineage = get_facets_method(*args)
@@ -153,5 +153,5 @@ class DefaultExtractor(BaseExtractor):
                 "This should not happen."
             )
         except Exception:
-            self.log.warning("OpenLineage provider method failed to extract data from provider. ")
+            self.log.warning("OpenLineage provider method failed to extract data from provider.")
         return None

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/extractors/manager.py

@@ -24,7 +24,11 @@ from airflow.providers.common.compat.openlineage.utils.utils import (
 )
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.extractors import BaseExtractor, OperatorLineage
-from airflow.providers.openlineage.extractors.base import DefaultExtractor
+from airflow.providers.openlineage.extractors.base import (
+    OL_METHOD_NAME_COMPLETE,
+    OL_METHOD_NAME_START,
+    DefaultExtractor,
+)
 from airflow.providers.openlineage.extractors.bash import BashExtractor
 from airflow.providers.openlineage.extractors.python import PythonExtractor
 from airflow.providers.openlineage.utils.utils import (
@@ -32,6 +36,7 @@ from airflow.providers.openlineage.utils.utils import (
     try_import_from_string,
 )
 from airflow.utils.log.logging_mixin import LoggingMixin
+from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     from openlineage.client.event_v2 import Dataset
@@ -87,7 +92,9 @@ class ExtractorManager(LoggingMixin):
     def add_extractor(self, operator_class: str, extractor: type[BaseExtractor]):
         self.extractors[operator_class] = extractor
 
-    def extract_metadata(self, dagrun, task, complete: bool = False, task_instance=None) -> OperatorLineage:
+    def extract_metadata(
+        self, dagrun, task, task_instance_state: TaskInstanceState, task_instance=None
+    ) -> OperatorLineage:
         extractor = self._get_extractor(task)
         task_info = (
             f"task_type={task.task_type} "
@@ -104,10 +111,15 @@ class ExtractorManager(LoggingMixin):
                     extractor.__class__.__name__,
                     str(task_info),
                 )
-                if complete:
-                    task_metadata = extractor.extract_on_complete(task_instance)
-                else:
+                if task_instance_state == TaskInstanceState.RUNNING:
                     task_metadata = extractor.extract()
+                elif task_instance_state == TaskInstanceState.FAILED:
+                    if callable(getattr(extractor, "extract_on_failure", None)):
+                        task_metadata = extractor.extract_on_failure(task_instance)
+                    else:
+                        task_metadata = extractor.extract_on_complete(task_instance)
+                else:
+                    task_metadata = extractor.extract_on_complete(task_instance)
 
                 self.log.debug(
                     "Found task metadata for operation %s: %s",
@@ -155,13 +167,9 @@ class ExtractorManager(LoggingMixin):
             return self.extractors[task.task_type]
 
         def method_exists(method_name):
-            method = getattr(task, method_name, None)
-            if method:
-                return callable(method)
+            return callable(getattr(task, method_name, None))
 
-        if method_exists("get_openlineage_facets_on_start") or method_exists(
-            "get_openlineage_facets_on_complete"
-        ):
+        if method_exists(OL_METHOD_NAME_START) or method_exists(OL_METHOD_NAME_COMPLETE):
             return self.default_extractor
         return None
 
@@ -191,7 +199,8 @@ class ExtractorManager(LoggingMixin):
             if d:
                 task_metadata.outputs.append(d)
 
-    def get_hook_lineage(self) -> tuple[list[Dataset], list[Dataset]] | None:
+    @staticmethod
+    def get_hook_lineage() -> tuple[list[Dataset], list[Dataset]] | None:
         try:
             from airflow.providers.common.compat.lineage.hook import (
                 get_hook_lineage_collector,

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/get_provider_info.py

@@ -27,8 +27,9 @@ def get_provider_info():
         "name": "OpenLineage Airflow",
         "description": "`OpenLineage <https://openlineage.io/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1741509355,
+        "source-date-epoch": 1742478177,
         "versions": [
+            "2.1.2",
             "2.1.1",
             "2.1.0",
             "2.0.0",

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/plugins/adapter.py

@@ -85,7 +85,7 @@ class OpenLineageAdapter(LoggingMixin):
             if config:
                 self.log.debug(
                     "OpenLineage configuration found. Transport type: `%s`",
-                    config.get("type", "no type provided"),
+                    config.get("transport", {}).get("type", "no type provided"),
                 )
                 self._client = OpenLineageClient(config=config)  # type: ignore[call-arg]
             else:

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/plugins/listener.py

@@ -69,13 +69,15 @@ def _get_try_number_success(val):
 
 def _executor_initializer():
     """
-    Initialize worker processes for the executor used for DagRun listener.
+    Initialize processes for the executor used with DAGRun listener's methods (on scheduler).
 
     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.
     """
+    # This initializer is used only on the scheduler
+    # We can configure_orm regardless of the Airflow version, as DB access is always allowed from scheduler.
     settings.configure_orm()
 
 
@@ -199,7 +201,9 @@ class OpenLineageListener:
             operator_name = task.task_type.lower()
 
             with Stats.timer(f"ol.extract.{event_type}.{operator_name}"):
-                task_metadata = self.extractor_manager.extract_metadata(dagrun, task)
+                task_metadata = self.extractor_manager.extract_metadata(
+                    dagrun=dagrun, task=task, task_instance_state=TaskInstanceState.RUNNING
+                )
 
             redacted_event = self.adapter.start_task(
                 run_id=task_uuid,
@@ -302,7 +306,10 @@ class OpenLineageListener:
 
             with Stats.timer(f"ol.extract.{event_type}.{operator_name}"):
                 task_metadata = self.extractor_manager.extract_metadata(
-                    dagrun, task, complete=True, task_instance=task_instance
+                    dagrun=dagrun,
+                    task=task,
+                    task_instance_state=TaskInstanceState.SUCCESS,
+                    task_instance=task_instance,
                 )
 
             redacted_event = self.adapter.complete_task(
@@ -423,7 +430,10 @@ class OpenLineageListener:
 
             with Stats.timer(f"ol.extract.{event_type}.{operator_name}"):
                 task_metadata = self.extractor_manager.extract_metadata(
-                    dagrun, task, complete=True, task_instance=task_instance
+                    dagrun=dagrun,
+                    task=task,
+                    task_instance_state=TaskInstanceState.FAILED,
+                    task_instance=task_instance,
                 )
 
             redacted_event = self.adapter.fail_task(
@@ -481,7 +491,8 @@ class OpenLineageListener:
             self.log.debug("Process with pid %s finished - parent", pid)
         else:
             setproctitle(getproctitle() + " - OpenLineage - " + callable_name)
-            configure_orm(disable_connection_pool=True)
+            if not AIRFLOW_V_3_0_PLUS:
+                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)

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/plugins/macros.py

@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter
 from airflow.providers.openlineage.utils.utils import get_job_name
+from airflow.providers.openlineage.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from airflow.models import TaskInstance
@@ -58,15 +59,25 @@ def lineage_run_id(task_instance: TaskInstance):
         For more information take a look at the guide:
         :ref:`howto/macros:openlineage`
     """
-    if hasattr(task_instance, "logical_date"):
-        logical_date = task_instance.logical_date
+    if AIRFLOW_V_3_0_PLUS:
+        context = task_instance.get_template_context()
+        if hasattr(task_instance, "dag_run"):
+            dag_run = task_instance.dag_run
+        elif hasattr(context, "dag_run"):
+            dag_run = context["dag_run"]
+        if hasattr(dag_run, "logical_date") and dag_run.logical_date:
+            date = dag_run.logical_date
+        else:
+            date = dag_run.run_after
+    elif hasattr(task_instance, "logical_date"):
+        date = task_instance.logical_date
     else:
-        logical_date = task_instance.execution_date
+        date = task_instance.execution_date
     return OpenLineageAdapter.build_task_instance_run_id(
         dag_id=task_instance.dag_id,
         task_id=task_instance.task_id,
         try_number=task_instance.try_number,
-        logical_date=logical_date,
+        logical_date=date,
         map_index=task_instance.map_index,
     )
 

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/utils/spark.py

@@ -53,35 +53,73 @@ def _get_parent_job_information_as_spark_properties(context: Context) -> dict:
 
 def _get_transport_information_as_spark_properties() -> dict:
     """Retrieve transport information as Spark properties."""
-    transport = get_openlineage_listener().adapter.get_or_create_openlineage_client().transport
-    if transport.kind != "http":
-        log.info(
-            "OpenLineage transport type `%s` does not support automatic "
-            "injection of OpenLineage transport information into Spark properties.",
-            transport.kind,
-        )
-        return {}
-
-    properties = {
-        "spark.openlineage.transport.type": transport.kind,
-        "spark.openlineage.transport.url": transport.url,
-        "spark.openlineage.transport.endpoint": transport.endpoint,
-        "spark.openlineage.transport.timeoutInMillis": str(
-            int(transport.timeout * 1000)  # convert to milliseconds, as required by Spark integration
-        ),
-    }
-    if transport.compression:
-        properties["spark.openlineage.transport.compression"] = str(transport.compression)
 
-    if hasattr(transport.config.auth, "api_key") and transport.config.auth.get_bearer():
-        properties["spark.openlineage.transport.auth.type"] = "api_key"
-        properties["spark.openlineage.transport.auth.apiKey"] = transport.config.auth.get_bearer()
+    def _get_transport_information(tp) -> dict:
+        properties = {
+            "type": tp.kind,
+            "url": tp.url,
+            "endpoint": tp.endpoint,
+            "timeoutInMillis": str(
+                int(tp.timeout) * 1000  # convert to milliseconds, as required by Spark integration
+            ),
+        }
+        if hasattr(tp, "compression") and tp.compression:
+            properties["compression"] = str(tp.compression)
+
+        if hasattr(tp.config.auth, "api_key") and tp.config.auth.get_bearer():
+            properties["auth.type"] = "api_key"
+            properties["auth.apiKey"] = tp.config.auth.get_bearer()
+
+        if hasattr(tp.config, "custom_headers") and tp.config.custom_headers:
+            for key, value in tp.config.custom_headers.items():
+                properties[f"headers.{key}"] = value
+        return properties
+
+    def _format_transport(props: dict, transport: dict, name: str | None):
+        for key, value in transport.items():
+            if name:
+                props[f"spark.openlineage.transport.transports.{name}.{key}"] = value
+            else:
+                props[f"spark.openlineage.transport.{key}"] = value
+        return props
 
-    if hasattr(transport.config, "custom_headers") and transport.config.custom_headers:
-        for key, value in transport.config.custom_headers.items():
-            properties[f"spark.openlineage.transport.headers.{key}"] = value
+    transport = get_openlineage_listener().adapter.get_or_create_openlineage_client().transport
 
-    return properties
+    if transport.kind == "composite":
+        http_transports = {}
+        for nested_transport in transport.transports:
+            if nested_transport.kind == "http":
+                http_transports[nested_transport.name] = _get_transport_information(nested_transport)
+            else:
+                name = nested_transport.name if hasattr(nested_transport, "name") else "no-name"
+                log.info(
+                    "OpenLineage transport type `%s` with name `%s` is not supported in composite transport.",
+                    nested_transport.kind,
+                    name,
+                )
+        if len(http_transports) == 0:
+            log.warning(
+                "OpenLineage transport type `composite` does not contain http transport. Skipping "
+                "injection of OpenLineage transport information into Spark properties.",
+            )
+            return {}
+        props = {
+            "spark.openlineage.transport.type": "composite",
+            "spark.openlineage.transport.continueOnFailure": str(transport.config.continue_on_failure),
+        }
+        for name, http_transport in http_transports.items():
+            props = _format_transport(props, http_transport, name)
+        return props
+
+    elif transport.kind == "http":
+        return _format_transport({}, _get_transport_information(transport), None)
+
+    log.info(
+        "OpenLineage transport type `%s` does not support automatic "
+        "injection of OpenLineage transport information into Spark properties.",
+        transport.kind,
+    )
+    return {}
 
 
 def _is_parent_job_information_present_in_spark_properties(properties: dict) -> bool:

{apache_airflow_providers_openlineage-2.1.1rc1 → apache_airflow_providers_openlineage-2.1.2rc1}/src/airflow/providers/openlineage/utils/utils.py

@@ -213,13 +213,7 @@ def is_ti_rescheduled_already(ti: TaskInstance, session=NEW_SESSION):
 
     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,
-            )
+            exists().where(TaskReschedule.ti_id == ti.id, TaskReschedule.try_number == ti.try_number)
         ).scalar()
         is True
     )
@@ -369,8 +363,19 @@ class DagRunInfo(InfoJsonEncodable):
         "run_id",
         "run_type",
         "start_date",
+        "end_date",
     ]
 
+    casts = {"duration": lambda dagrun: DagRunInfo.duration(dagrun)}
+
+    @classmethod
+    def duration(cls, dagrun: DagRun) -> float | None:
+        if not getattr(dagrun, "end_date", None) or not isinstance(dagrun.end_date, datetime.datetime):
+            return None
+        if not getattr(dagrun, "start_date", None) or not isinstance(dagrun.start_date, datetime.datetime):
+            return None
+        return (dagrun.end_date - dagrun.start_date).total_seconds()
+
 
 class TaskInstanceInfo(InfoJsonEncodable):
     """Defines encoding TaskInstance object to JSON."""