fix: use opened compute for telemetry SPOG and prune closed coordinator states

Author: saishreeeee

dbt/adapters/databricks/connections.py

@@ -522,7 +522,9 @@ def connect() -> DatabricksHandle:
                         databricks_connection.http_path
                     )
 
-                    telemetry_hooks.on_connection_open(creds, credentials_manager)
+                    telemetry_hooks.on_connection_open(
+                        creds, credentials_manager, databricks_connection.http_path
+                    )
                     return conn
                 else:
                     raise DbtDatabaseError("Failed to create connection")

dbt/adapters/databricks/telemetry/builder.py

@@ -186,6 +186,18 @@ def build_invocation_config(config: Any) -> models.InvocationConfig:
     )
 
 
+def _profile_http_paths(creds: DatabricksCredentials) -> list[str]:
+    paths: list[str] = []
+    default = getattr(creds, "http_path", None)
+    if default:
+        paths.append(default)
+    for cfg in (getattr(creds, "compute", None) or {}).values():
+        path = cfg.get("http_path") if cfg else None
+        if path:
+            paths.append(path)
+    return paths
+
+
 def build_connection_config(creds: DatabricksCredentials) -> models.ConnectionConfig:
     http_path = getattr(creds, "http_path", None)
     connection_parameters = getattr(creds, "connection_parameters", None) or {}
@@ -194,7 +206,9 @@ def build_connection_config(creds: DatabricksCredentials) -> models.ConnectionCo
         configured_auth_family=classify_auth_family(creds),
         named_compute_count=len(getattr(creds, "compute", None) or {}),
         # Parse only the `o` parameter; discard its value.
-        spog_routing_configured=extract_workspace_id(http_path) is not None,
+        spog_routing_configured=any(
+            extract_workspace_id(path) is not None for path in _profile_http_paths(creds)
+        ),
         use_kernel=bool(connection_parameters.get("use_kernel")),
     )
 

dbt/adapters/databricks/telemetry/coordinator.py

@@ -119,6 +119,11 @@ def set_transport(self, invocation_id: str, transport: Transport) -> None:
 
     def mark_start(self, invocation_id: str) -> None:
         with self._lock:
+            self._states = {
+                key: state
+                for key, state in self._states.items()
+                if key == invocation_id or not state.closed
+            }
             self._state(invocation_id)
 
     def is_closed(self, invocation_id: str) -> bool:

dbt/adapters/databricks/telemetry/hooks.py

@@ -2,6 +2,7 @@
 from typing import Any, Optional
 
 from dbt.adapters.databricks.credentials import DatabricksCredentials
+from dbt.adapters.databricks.spog.extract import extract_workspace_id
 from dbt.adapters.databricks.telemetry import builder, listener
 from dbt.adapters.databricks.telemetry.config import (
     has_reusable_transport,
@@ -73,6 +74,7 @@ def on_post_parse(adapter: Any, manifest: Any) -> None:
 def on_connection_open(
     credentials: Optional[DatabricksCredentials],
     credentials_manager: Optional[Any],
+    http_path: Optional[str] = None,
 ) -> None:
     try:
         if (
@@ -84,10 +86,13 @@ def on_connection_open(
         invocation_id = _current_invocation_id()
         if not invocation_id:
             return
+        workspace_id = extract_workspace_id(http_path)
+        if workspace_id is None:
+            workspace_id = getattr(credentials_manager, "workspace_id", None)
         transport = Transport(
             host=getattr(credentials_manager, "host", None),
             header_factory=getattr(credentials_manager, "header_factory", None),
-            workspace_id=getattr(credentials_manager, "workspace_id", None),
+            workspace_id=workspace_id,
         )
         coordinator().set_transport(invocation_id, transport)
     except Exception:  # pragma: no cover - best-effort

tests/unit/telemetry/test_builder.py

@@ -162,6 +162,16 @@ def test_spog_parameter_is_parsed_not_substring_matched(self):
         )
         assert cc.spog_routing_configured is False
 
+    def test_named_compute_o_parameter_sets_spog_flag(self):
+        cc = builder.build_connection_config(
+            _creds(
+                token="dapi",
+                http_path="/sql/1.0/warehouses/default",
+                compute={"named": {"http_path": "/sql/1.0/warehouses/named?o=42"}},
+            )
+        )
+        assert cc.spog_routing_configured is True
+
 
 class TestAggregateManifest:
     def _manifest(self):

tests/unit/telemetry/test_coordinator.py

@@ -187,6 +187,20 @@ def start(self):
 
         assert len(capture.calls) == 2
 
+    def test_mark_start_prunes_closed_invocations(self):
+        c = coord_mod.Coordinator()
+        c.mark_start("inv-1")
+        c.close("inv-1")
+        c.mark_start("inv-2")
+
+        assert "inv-1" not in c._states
+        assert "inv-2" in c._states
+        assert c.needs_post_parse("inv-2") is True
+
+        c.close("inv-1")
+        assert c._states["inv-1"].closed
+        assert not c._states["inv-2"].closed
+
 
 class TestPostRun:
     def _entry(self, call):

tests/unit/telemetry/test_hooks.py

@@ -76,3 +76,33 @@ def test_run_end_exception_finalizes_stored_invocation_not_current_global(monkey
     assert build.call_args.args[0] == "inv-1"
     coord.set_post_run.assert_called_once_with("inv-1", "log")
     coord.close.assert_called_once_with("inv-1")
+
+
+def test_connection_open_uses_opened_path_workspace_id(monkeypatch):
+    coord = Mock()
+    monkeypatch.setattr(hooks, "coordinator", lambda: coord)
+    monkeypatch.setattr(hooks, "_current_invocation_id", lambda: "inv-1")
+    monkeypatch.setattr(hooks, "is_enabled_for_invocation", lambda _: True)
+    monkeypatch.setattr(hooks, "has_reusable_transport", lambda _: True)
+    manager = SimpleNamespace(host="https://h", header_factory=lambda: {}, workspace_id=None)
+
+    hooks.on_connection_open(
+        SimpleNamespace(), manager, "/sql/1.0/warehouses/named?o=42"
+    )
+
+    transport = coord.set_transport.call_args.args[1]
+    assert transport.workspace_id == "42"
+
+
+def test_connection_open_falls_back_to_manager_workspace_id(monkeypatch):
+    coord = Mock()
+    monkeypatch.setattr(hooks, "coordinator", lambda: coord)
+    monkeypatch.setattr(hooks, "_current_invocation_id", lambda: "inv-1")
+    monkeypatch.setattr(hooks, "is_enabled_for_invocation", lambda _: True)
+    monkeypatch.setattr(hooks, "has_reusable_transport", lambda _: True)
+    manager = SimpleNamespace(host="https://h", header_factory=lambda: {}, workspace_id="7")
+
+    hooks.on_connection_open(SimpleNamespace(), manager, "/sql/1.0/warehouses/default")
+
+    transport = coord.set_transport.call_args.args[1]
+    assert transport.workspace_id == "7"