fix: bind telemetry cleanup to the adapter invocation and claim queued sends

Repeated dbtRunner invokes reset the process-global invocation ID before leftover adapter cleanup, and close() can drop a POST_PARSE event that was already scheduled.

Author: saishreeeee

dbt/adapters/databricks/telemetry/coordinator.py

@@ -91,13 +91,15 @@ def needs_post_parse(self, invocation_id: str) -> bool:
             return state is None or (not state.closed and state.post_parse is None)
 
     def send_if_ready(self, invocation_id: str) -> None:
-        # Keep connection-open off the network path.
+        # Keep connection-open off the network path. Claim send inputs under
+        # the lock so a later close() cannot cancel already-scheduled work.
         with self._lock:
-            if not self._ready_to_send(self._states.get(invocation_id)):
+            claimed = self._claim_send(self._states.get(invocation_id))
+            if claimed is None:
                 return
         thread = threading.Thread(
-            target=self._drain,
-            args=(invocation_id,),
+            target=self._send_claimed,
+            args=(invocation_id, *claimed),
             name="dbt-telemetry-send",
             daemon=True,
         )
@@ -114,6 +116,26 @@ def _ready_to_send(self, state: Optional[_InvocationState]) -> bool:
             return False
         return state.post_parse is not None and not state.post_parse_sent
 
+    def _claim_send(
+        self, state: Optional[_InvocationState]
+    ) -> Optional[tuple[Optional[str], TelemetryLog, str, HeaderFactory, Optional[Any]]]:
+        if not self._ready_to_send(state) or state is None or state.transport is None:
+            return None
+        if state.post_parse is None or state.transport.header_factory is None:
+            return None
+        transport = state.transport
+        payload = state.post_parse
+        event_id = state.post_parse_event_id
+        state.post_parse_sent = True
+        state.sending = True
+        return (
+            transport.host,
+            payload,
+            event_id,
+            transport.header_factory,
+            transport.workspace_id,
+        )
+
     def flush(self, timeout: Optional[float] = None) -> None:
         if timeout is None:
             timeout = float(client._TIMEOUT_SECONDS * 2)
@@ -129,26 +151,16 @@ def flush(self, timeout: Optional[float] = None) -> None:
                 return
             pending[0].join(remaining)
 
-    def _drain(self, invocation_id: str) -> None:
-        with self._lock:
-            state = self._states.get(invocation_id)
-            if state is None or state.closed or state.sending:
-                return
-            transport = state.transport
-            if transport is None or transport.header_factory is None:
-                return
-            if state.post_parse is None or state.post_parse_sent:
-                return
-            payload = state.post_parse
-            event_id = state.post_parse_event_id
-            state.post_parse_sent = True
-            state.sending = True
-            host = transport.host
-            header_factory = transport.header_factory
-            workspace_id = transport.workspace_id
-
+    def _send_claimed(
+        self,
+        invocation_id: str,
+        host: Optional[str],
+        payload: TelemetryLog,
+        event_id: str,
+        header_factory: HeaderFactory,
+        workspace_id: Optional[Any],
+    ) -> None:
         self._send(host, payload, event_id, header_factory, workspace_id)
-
         with self._lock:
             state = self._states.get(invocation_id)
             if state is None or state.closed:

dbt/adapters/databricks/telemetry/hooks.py

@@ -8,6 +8,10 @@
 )
 from dbt.adapters.databricks.telemetry.coordinator import Transport, coordinator
 
+# Bound at adapter init. dbt-core resets the process-global invocation ID
+# before leftover adapter cleanup, so teardown must not read the live ID.
+_INVOCATION_ID_ATTR = "_dbt_telemetry_invocation_id"
+
 
 def _current_invocation_id() -> Optional[str]:
     try:
@@ -19,6 +23,11 @@ def _current_invocation_id() -> Optional[str]:
         return None
 
 
+def _stored_invocation_id(adapter: Any) -> Optional[str]:
+    invocation_id = getattr(adapter, _INVOCATION_ID_ATTR, None)
+    return str(invocation_id) if invocation_id else None
+
+
 def on_adapter_init(adapter: Any) -> None:
     try:
         creds = getattr(getattr(adapter, "config", None), "credentials", None)
@@ -27,6 +36,7 @@ def on_adapter_init(adapter: Any) -> None:
         invocation_id = _current_invocation_id()
         if not invocation_id:
             return
+        setattr(adapter, _INVOCATION_ID_ATTR, invocation_id)
         coordinator().mark_start(invocation_id)
     except Exception:  # pragma: no cover - best-effort
         return
@@ -85,7 +95,7 @@ def on_run_end(adapter: Any) -> None:
         creds = getattr(config, "credentials", None)
         if not isinstance(creds, DatabricksCredentials) or not is_enabled_for_invocation(creds):
             return
-        invocation_id = _current_invocation_id()
+        invocation_id = _stored_invocation_id(adapter)
         if not invocation_id:
             return
         coordinator().close(invocation_id)

tests/unit/telemetry/test_coordinator.py

@@ -121,3 +121,30 @@ def test_closed_invocation_rejects_late_callbacks(self, monkeypatch):
         c.set_transport("inv-1", _transport())
         c.flush()
         assert capture.calls == []
+
+    def test_close_does_not_cancel_already_queued_send(self, monkeypatch):
+        pending = []
+
+        class DelayedThread:
+            def __init__(self, target, args, name, daemon):
+                self.target = target
+                self.args = args
+
+            def is_alive(self):
+                return False
+
+            def start(self):
+                pending.append(self)
+
+        capture = _Capture()
+        monkeypatch.setattr(coord_mod.threading, "Thread", DelayedThread)
+        monkeypatch.setattr(coord_mod.client, "send", capture)
+        c = coord_mod.Coordinator()
+        c.set_post_parse("inv-1", _log())
+        c.set_transport("inv-1", _transport())
+
+        assert len(pending) == 1
+        c.close("inv-1")
+        pending[0].target(*pending[0].args)
+
+        assert len(capture.calls) == 1

tests/unit/telemetry/test_hooks.py

@@ -2,6 +2,7 @@
 from unittest.mock import Mock
 
 from dbt.adapters.databricks.telemetry import hooks
+from dbt.adapters.databricks.telemetry.coordinator import Coordinator
 
 
 def test_post_parse_is_not_rebuilt(monkeypatch):
@@ -24,12 +25,51 @@ def test_post_parse_is_not_rebuilt(monkeypatch):
 def test_run_end_closes_without_waiting(monkeypatch):
     coord = Mock()
     monkeypatch.setattr(hooks, "coordinator", lambda: coord)
-    monkeypatch.setattr(hooks, "_current_invocation_id", lambda: "inv-1")
     monkeypatch.setattr(hooks, "DatabricksCredentials", object)
     monkeypatch.setattr(hooks, "is_enabled_for_invocation", lambda _: True)
-    adapter = SimpleNamespace(config=SimpleNamespace(credentials=SimpleNamespace()))
+    adapter = SimpleNamespace(
+        config=SimpleNamespace(credentials=SimpleNamespace()),
+        _dbt_telemetry_invocation_id="inv-1",
+    )
 
     hooks.on_run_end(adapter)
 
     coord.flush.assert_not_called()
     coord.close.assert_called_once_with("inv-1")
+
+
+def test_run_end_closes_stored_invocation_not_current_global(monkeypatch):
+    coord = Mock()
+    monkeypatch.setattr(hooks, "coordinator", lambda: coord)
+    monkeypatch.setattr(hooks, "_current_invocation_id", lambda: "inv-2")
+    monkeypatch.setattr(hooks, "DatabricksCredentials", object)
+    monkeypatch.setattr(hooks, "is_enabled_for_invocation", lambda _: True)
+    adapter = SimpleNamespace(
+        config=SimpleNamespace(credentials=SimpleNamespace()),
+        _dbt_telemetry_invocation_id="inv-1",
+    )
+
+    hooks.on_run_end(adapter)
+
+    coord.close.assert_called_once_with("inv-1")
+
+
+def test_stale_cleanup_does_not_tombstone_next_invocation(monkeypatch):
+    coord = Coordinator()
+    current = ["inv-1"]
+    monkeypatch.setattr(hooks, "coordinator", lambda: coord)
+    monkeypatch.setattr(hooks, "DatabricksCredentials", object)
+    monkeypatch.setattr(hooks, "is_enabled_for_invocation", lambda _: True)
+    monkeypatch.setattr(hooks, "_current_invocation_id", lambda: current[0])
+    leftover = SimpleNamespace(config=SimpleNamespace(credentials=object()))
+
+    hooks.on_adapter_init(leftover)
+    hooks.on_run_end(leftover)
+    current[0] = "inv-2"
+    hooks.on_run_end(leftover)
+    next_adapter = SimpleNamespace(config=SimpleNamespace(credentials=object()))
+    hooks.on_adapter_init(next_adapter)
+
+    assert leftover._dbt_telemetry_invocation_id == "inv-1"
+    assert next_adapter._dbt_telemetry_invocation_id == "inv-2"
+    assert coord.needs_post_parse("inv-2") is True