fix(adapters/databricks/connections.py): fix capability detection for named compute

Author: trouze

dbt/adapters/databricks/connections.py

@@ -208,6 +208,23 @@ def _cache_dbr_capabilities(cls, creds: DatabricksCredentials, http_path: str) -
                 is_sql_warehouse=not is_cluster,
             )
 
+    @classmethod
+    def _try_cache_dbr_capabilities(cls, creds: DatabricksCredentials, http_path: str) -> None:
+        """Like _cache_dbr_capabilities, but only writes to the cache when the version query
+        actually succeeds. This prevents a failed eager lookup (e.g. credentials_manager not yet
+        set, cluster still spinning up) from storing a None-version entry that the idempotency
+        guard in open() would later treat as authoritative, causing all capability checks to
+        return False.
+        """
+        if http_path not in cls._dbr_capabilities_cache:
+            is_cluster = is_cluster_http_path(http_path, creds.cluster_id)
+            dbr_version = cls._query_dbr_version(creds, http_path)
+            if dbr_version is not None:
+                cls._dbr_capabilities_cache[http_path] = DBRCapabilities(
+                    dbr_version=dbr_version,
+                    is_sql_warehouse=not is_cluster,
+                )
+
     def cancel_open(self) -> list[str]:
         cancelled = super().cancel_open()
         logger.info("Cancelling open python jobs")
@@ -288,7 +305,7 @@ def _create_fresh_connection(
         conn.http_path = QueryConfigUtils.get_http_path(query_header_context, creds)
         conn.thread_identifier = cast(tuple[int, int], self.get_thread_identifier())
         conn._query_header_context = query_header_context
-        self._cache_dbr_capabilities(creds, conn.http_path)
+        self._try_cache_dbr_capabilities(creds, conn.http_path)
         conn.capabilities = self._get_capabilities_for_http_path(conn.http_path)
         conn.handle = LazyHandle(self.open)
 

tests/unit/test_connection_manager.py

@@ -124,51 +124,103 @@ def connection_manager(self):
         DatabricksConnectionManager._dbr_capabilities_cache = {}
         return mgr
 
-    @patch.object(DatabricksConnectionManager, "_cache_dbr_capabilities")
+    @patch.object(DatabricksConnectionManager, "_try_cache_dbr_capabilities")
     def test_fresh_connection_caches_capabilities_before_setting_them(
-        self, mock_cache_dbr, connection_manager
+        self, mock_try_cache, connection_manager
     ):
-        """_create_fresh_connection must call _cache_dbr_capabilities before reading the cache."""
+        """_create_fresh_connection must call _try_cache_dbr_capabilities before reading the
+        cache."""
         expected_caps = DBRCapabilities(dbr_version=(15, 4), is_sql_warehouse=False)
 
         def populate_cache(creds, http_path):
             DatabricksConnectionManager._dbr_capabilities_cache[http_path] = expected_caps
 
-        mock_cache_dbr.side_effect = populate_cache
+        mock_try_cache.side_effect = populate_cache
 
         ctx = QueryContextWrapper(compute_name="alternate_compute")
         conn = connection_manager._create_fresh_connection("test_model", ctx)
 
-        mock_cache_dbr.assert_called_once_with(
+        mock_try_cache.assert_called_once_with(
             connection_manager.profile.credentials, self.NAMED_COMPUTE_HTTP_PATH
         )
         assert conn.capabilities == expected_caps
         assert conn.capabilities.has_capability(DBRCapability.ICEBERG)
 
-    @patch.object(DatabricksConnectionManager, "_cache_dbr_capabilities")
-    def test_fresh_connection_without_eager_cache_gets_empty_defaults(
-        self, mock_cache_dbr, connection_manager
+    @patch.object(DatabricksConnectionManager, "_try_cache_dbr_capabilities")
+    def test_fresh_connection_falls_back_to_empty_defaults_when_eager_cache_is_noop(
+        self, mock_try_cache, connection_manager
     ):
-        """Without the eager cache call, named compute gets empty (broken) capabilities.
-
-        This test documents the pre-fix behavior: if _cache_dbr_capabilities is a no-op,
-        the connection falls back to DBRCapabilities() which has all capabilities disabled.
-        """
+        """When the eager try-cache is a no-op (e.g. version query not yet possible), the
+        connection gets default capabilities. open() will populate the cache correctly later."""
         ctx = QueryContextWrapper(compute_name="alternate_compute")
         conn = connection_manager._create_fresh_connection("test_model", ctx)
 
         assert conn.capabilities.dbr_version is None
         assert not conn.capabilities.is_sql_warehouse
         assert not conn.capabilities.has_capability(DBRCapability.ICEBERG)
 
-    @patch.object(DatabricksConnectionManager, "_cache_dbr_capabilities")
-    def test_default_compute_cache_is_idempotent(self, mock_cache_dbr, connection_manager):
+    @patch.object(DatabricksConnectionManager, "_try_cache_dbr_capabilities")
+    def test_default_compute_cache_is_idempotent(self, mock_try_cache, connection_manager):
         """For default compute with a warm cache, the eager call is a no-op."""
         warm_caps = DBRCapabilities(dbr_version=(15, 4), is_sql_warehouse=False)
         DatabricksConnectionManager._dbr_capabilities_cache[self.DEFAULT_HTTP_PATH] = warm_caps
 
         ctx = QueryContextWrapper()
         conn = connection_manager._create_fresh_connection("test_model", ctx)
 
-        mock_cache_dbr.assert_called_once()
+        mock_try_cache.assert_called_once()
         assert conn.capabilities == warm_caps
+
+
+class TestTryCacheDbr:
+    """Unit tests for _try_cache_dbr_capabilities."""
+
+    HTTP_PATH = "sql/protocolv1/o/1234567890123456/cluster-abc"
+
+    @pytest.fixture(autouse=True)
+    def clear_cache(self):
+        DatabricksConnectionManager._dbr_capabilities_cache = {}
+        yield
+        DatabricksConnectionManager._dbr_capabilities_cache = {}
+
+    @patch.object(DatabricksConnectionManager, "_query_dbr_version", return_value=None)
+    def test_does_not_write_to_cache_when_version_is_none(self, mock_query):
+        """When the version query returns None, the cache must not be written.
+
+        This prevents a poisoned None entry from blocking the authoritative write in open().
+        """
+        creds = Mock(spec=DatabricksCredentials)
+        creds.cluster_id = None
+
+        DatabricksConnectionManager._try_cache_dbr_capabilities(creds, self.HTTP_PATH)
+
+        assert self.HTTP_PATH not in DatabricksConnectionManager._dbr_capabilities_cache
+        mock_query.assert_called_once_with(creds, self.HTTP_PATH)
+
+    @patch.object(DatabricksConnectionManager, "_query_dbr_version", return_value=(15, 4))
+    def test_writes_to_cache_when_version_is_known(self, mock_query):
+        """When the version query succeeds, capabilities are cached correctly."""
+        creds = Mock(spec=DatabricksCredentials)
+        creds.cluster_id = None
+
+        DatabricksConnectionManager._try_cache_dbr_capabilities(creds, self.HTTP_PATH)
+
+        mock_query.assert_called_once_with(creds, self.HTTP_PATH)
+        caps = DatabricksConnectionManager._dbr_capabilities_cache.get(self.HTTP_PATH)
+        assert caps is not None
+        assert caps.dbr_version == (15, 4)
+        assert not caps.is_sql_warehouse
+        assert caps.has_capability(DBRCapability.ICEBERG)
+
+    @patch.object(DatabricksConnectionManager, "_query_dbr_version", return_value=(15, 4))
+    def test_skips_write_when_already_cached(self, mock_query):
+        """If the path is already in cache, the version query is never made."""
+        creds = Mock(spec=DatabricksCredentials)
+        creds.cluster_id = None
+        existing = DBRCapabilities(dbr_version=(14, 3), is_sql_warehouse=False)
+        DatabricksConnectionManager._dbr_capabilities_cache[self.HTTP_PATH] = existing
+
+        DatabricksConnectionManager._try_cache_dbr_capabilities(creds, self.HTTP_PATH)
+
+        mock_query.assert_not_called()
+        assert DatabricksConnectionManager._dbr_capabilities_cache[self.HTTP_PATH] is existing