Fix OpenLineage team_name access and other failing paths

providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py

@@ -21,7 +21,6 @@
 import sys
 import threading
 from concurrent.futures import ProcessPoolExecutor
-from concurrent.futures.process import BrokenProcessPool
 from datetime import datetime
 from functools import cache
 from typing import TYPE_CHECKING
@@ -138,8 +137,12 @@ def _run_adapter_method(method_name: str, /, *args, **kwargs):
     background worker threads (e.g. the ``datadog`` transport, which always starts an async
     HTTP worker thread) are never closed, so this leaks one thread per event and steadily
     consumes scheduler CPU and memory until restart.
+
+    Returns nothing: the emitted event is only of use inside the pool worker, and pickling it
+    back to the parent would turn a redacted event the pickler chokes on into a spurious
+    "failed to submit" warning for an emission that actually succeeded.
     """
-    return getattr(_get_process_adapter(), method_name)(*args, **kwargs)
+    getattr(_get_process_adapter(), method_name)(*args, **kwargs)
 
 
 def _emit_manual_state_change_event(adapter_method_name: str, stats_key: str, **kwargs):
@@ -149,11 +152,11 @@ def _emit_manual_state_change_event(adapter_method_name: str, stats_key: str, **
     Module-level so it is picklable across the ProcessPoolExecutor boundary used by
     `_on_task_instance_manual_state_change` for scheduler-side "task state changed
     externally" emissions. The method is resolved on the per-process adapter so the
-    pool worker reuses one client across events (see ``_run_adapter_method``).
+    pool worker reuses one client across events (see ``_run_adapter_method``), and nothing is
+    returned so an unpicklable event cannot fail the future after a successful emission.
     """
     event = getattr(_get_process_adapter(), adapter_method_name)(**kwargs)
     Stats.gauge(stats_key, len(Serde.to_json(event).encode("utf-8")))
-    return event
 
 
 class OpenLineageListener:
@@ -1058,14 +1061,21 @@ def _fork_execute(self, callable, callable_name: str):
                 self._terminate_with_wait(process)
             self.log.debug("Process with pid %s finished - parent", pid)
         else:
-            setproctitle(getproctitle() + " - OpenLineage - " + callable_name)
-            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())
+            # Everything the child does lives in this try, and the exit lives in its finally: a child
+            # that returns instead of exiting becomes a second task runner. It would unwind into the
+            # hook's caller and keep executing task-runner code while sharing the supervisor socket
+            # with the real one, duplicating state writes and interleaving bytes on the channel.
             try:
+                setproctitle(getproctitle() + " - OpenLineage - " + callable_name)
+                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)
-            except Exception:
+            except BaseException:
+                # BaseException, not Exception: a SIGINT delivered to the process group reaches this
+                # child as KeyboardInterrupt, and Airflow's own AirflowTaskTimeout / TaskDeferred
+                # also derive from BaseException.
                 self.log.warning(
                     "OpenLineage %s process failed. This has no impact on actual task execution status.",
                     callable_name,
@@ -1076,8 +1086,13 @@ def _fork_execute(self, callable, callable_name: str):
                 # logging so buffered records (including any warnings above) are flushed
                 # before the process exits. Without this, the final log lines are silently
                 # dropped, making failures invisible.
-                logging.shutdown()
-            os._exit(0)
+                # Nest os._exit in its own finally so that a raising logging.shutdown()
+                # (e.g. a remote handler whose close() throws) cannot skip the exit and
+                # unwind back into the task runner as a second process.
+                try:
+                    logging.shutdown()
+                finally:
+                    os._exit(0)
 
     @property
     def executor(self) -> ProcessPoolExecutor:
@@ -1095,8 +1110,26 @@ def on_starting(self, component) -> None:
     @hookimpl
     def before_stopping(self, component) -> None:
         self.log.debug("before_stopping: %s", component.__class__.__name__)
-        with timeout(30):
-            self.executor.shutdown(wait=True)
+        if self._executor is None:
+            # Do not build a pool just to tear it down -- on the task runner this hook fires at the
+            # end of every task, where no pool was ever needed.
+            return
+
+        # Detach before shutting down so a later event rebuilds a fresh pool. Left attached, every
+        # subsequent submission would raise "cannot schedule new futures after shutdown" and drop
+        # its event for the remaining lifetime of the process.
+        executor, self._executor = self._executor, None
+        try:
+            with timeout(30):
+                executor.shutdown(wait=True)
+        except BaseException:
+            # `timeout` is SIGALRM-based: it raises AirflowTaskTimeout when shutdown overruns, and
+            # ValueError when called off the main thread. Neither may escape a listener hook --
+            # and BaseException is required, not Exception, because AirflowTaskTimeout derives from
+            # BaseException so that user code cannot swallow it. Every hook call site guards with
+            # `except Exception`, so letting it through would reach the caller.
+            self.log.warning("OpenLineage executor did not shut down cleanly.", exc_info=True)
+            executor.shutdown(wait=False)
 
     @hookimpl
     def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None:
@@ -1258,9 +1291,14 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None:
     def submit_callable(self, callable, *args, **kwargs):
         try:
             fut = self.executor.submit(callable, *args, **kwargs)
-        except BrokenProcessPool:
-            self.log.warning("ProcessPoolExecutor is broken; recreating and retrying submission.")
-            self._executor.shutdown(wait=False)
+        except RuntimeError:
+            # BrokenProcessPool subclasses RuntimeError, so this also covers a pool that was already
+            # shut down ("cannot schedule new futures after shutdown"). The retry is deliberately
+            # unguarded: every caller wraps this in `except BaseException`, and a second consecutive
+            # failure deserves to surface in their warning rather than be swallowed here.
+            self.log.warning("ProcessPoolExecutor is unusable; recreating and retrying submission.")
+            if self._executor is not None:
+                self._executor.shutdown(wait=False)
             self._executor = None
             fut = self.executor.submit(callable, *args, **kwargs)
         fut.add_done_callback(self.log_submit_error)

providers/openlineage/src/airflow/providers/openlineage/utils/utils.py

@@ -1045,15 +1045,18 @@ def dag_version_info(cls, dagrun: DagRun, key: str) -> str | int | None:
         dag_versions = safe_getattr(dagrun, "dag_versions", [])
         if not dag_versions:
             return None
+        # The DagVersion rows themselves are lazy-loaded, so reading their columns can still hit
+        # a detached session even though fetching the list above succeeded.
         current_version = dag_versions[-1]
         if key == "bundle_name":
-            return current_version.bundle_name
+            return safe_getattr(current_version, "bundle_name")
         if key == "bundle_version":
-            return current_version.bundle_version
+            return safe_getattr(current_version, "bundle_version")
         if key == "version_id":
-            return str(current_version.id)
+            version_id = safe_getattr(current_version, "id")
+            return str(version_id) if version_id is not None else None
         if key == "version_number":
-            return current_version.version_number
+            return safe_getattr(current_version, "version_number")
         raise ValueError(f"Unsupported key: {key}`")
 
     @classmethod
@@ -1062,13 +1065,30 @@ def team_name(cls, dagrun: DagRun) -> str | None:
         if not AIRFLOW_V_3_3_PLUS or not airflow_conf.getboolean("core", "multi_team", fallback=False):
             return None
 
-        from airflow.models.dagbundle import DagBundleModel
+        # The Execution API delivers the team name on the DagRun payload it sends to the task
+        # runner, so task events resolve it from there rather than through the bundle lookup below,
+        # which needs a metadata DB session the task runner does not have.
+        # `hasattr` rather than a None check -- a team-less run legitimately carries None, while
+        # the scheduler's ORM DagRun has no such attribute at all.
+        if hasattr(dagrun, "team_name"):
+            return dagrun.team_name
 
-        bundle_name = cls.dag_version_info(dagrun, "bundle_name")
-        if not isinstance(bundle_name, str):
-            return None
+        try:
+            bundle_name = cls.dag_version_info(dagrun, "bundle_name")
+            if not isinstance(bundle_name, str):
+                return None
 
-        return DagBundleModel.get_team_name(bundle_name)
+            from airflow.models.dagbundle import DagBundleModel
+
+            return DagBundleModel.get_team_name(bundle_name)
+        except Exception as e:
+            log.warning(
+                "OpenLineage failed to resolve the team name for dag `%s`: %s.",
+                safe_getattr(dagrun, "dag_id"),
+                e,
+            )
+            log.debug("Exception details:", exc_info=True)
+            return None
 
 
 class TaskInstanceInfo(InfoJsonEncodable):
@@ -1387,23 +1407,28 @@ def get_airflow_job_facet(dag_run: DagRun) -> dict[str, AirflowJobFacet]:
 def get_airflow_state_run_facet(
     dag_id: str, run_id: str, task_ids: list[str], dag_run_state: DagRunState
 ) -> dict[str, AirflowStateRunFacet]:
-    tis = DagRun.fetch_task_instances(dag_id=dag_id, run_id=run_id, task_ids=task_ids)
+    try:
+        tis = DagRun.fetch_task_instances(dag_id=dag_id, run_id=run_id, task_ids=task_ids)
 
-    def get_task_duration(ti):
-        if ti.duration is not None:
-            return ti.duration
-        if ti.end_date is not None and ti.start_date is not None:
-            return (ti.end_date - ti.start_date).total_seconds()
-        # Fallback to 0.0 for tasks with missing timestamps (e.g., skipped/terminated tasks)
-        return 0.0
+        def get_task_duration(ti):
+            if ti.duration is not None:
+                return ti.duration
+            if ti.end_date is not None and ti.start_date is not None:
+                return (ti.end_date - ti.start_date).total_seconds()
+            # Fallback to 0.0 for tasks with missing timestamps (e.g., skipped/terminated tasks)
+            return 0.0
 
-    return {
-        "airflowState": AirflowStateRunFacet(
-            dagRunState=dag_run_state,
-            tasksState={ti.task_id: ti.state for ti in tis},
-            tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis},
-        )
-    }
+        return {
+            "airflowState": AirflowStateRunFacet(
+                dagRunState=dag_run_state,
+                tasksState={ti.task_id: ti.state for ti in tis},
+                tasksDuration={ti.task_id: get_task_duration(ti) for ti in tis},
+            )
+        }
+    except Exception as e:
+        log.warning("Failed to build AirflowStateRunFacet for DagRun %s/%s: %s.", dag_id, run_id, e)
+        log.debug("Exception details:", exc_info=True)
+        return {}
 
 
 def is_dag_run_asset_triggered(