neutral state catalog

Author: fresioAS

sqlmesh/core/engine_adapter/fabric.py

@@ -52,6 +52,23 @@ def _target_catalog(self) -> t.Optional[str]:
     def _target_catalog(self, value: t.Optional[str]) -> None:
         self._connection_pool.set_attribute("target_catalog", value)
 
+    def _normalize_catalog(
+        self, catalog_name: t.Optional[str]
+    ) -> t.Optional[str]:
+        if not catalog_name:
+            return None
+
+        default_catalog = (
+            self._default_catalog or self._extra_config.get("database")
+        )
+        if default_catalog and catalog_name == default_catalog:
+            return None
+
+        return catalog_name
+
+    def _catalog_state_label(self, catalog_name: t.Optional[str]) -> str:
+        return catalog_name or "<default>"
+
     @property
     def api_client(self) -> FabricHttpClient:
         # the requests Session is not guaranteed to be threadsafe
@@ -109,10 +126,10 @@ def _drop_catalog(self, catalog_name: exp.Identifier) -> None:
             self.close()
 
     def get_current_catalog(self) -> t.Optional[str]:
-        """Return the adapter-managed catalog for Fabric's stateless sessions."""
-        return self._target_catalog or self._extra_config.get("database")
+        """Return the explicit Fabric catalog target for the current thread."""
+        return self._normalize_catalog(self._target_catalog)
 
-    def set_current_catalog(self, catalog_name: str) -> None:
+    def set_current_catalog(self, catalog_name: t.Optional[str]) -> None:
         """
         Set the current catalog for Microsoft Fabric connections.
 
@@ -121,7 +138,8 @@ def set_current_catalog(self, catalog_name: str) -> None:
         recreate them with the new catalog in the connection configuration.
 
         Args:
-            catalog_name: The name of the catalog (warehouse) to switch to
+            catalog_name: The name of the catalog (warehouse) to switch to.
+                The configured default catalog is treated as the neutral state.
 
         Note:
             Fabric doesn't support catalog switching via USE statements because each
@@ -132,13 +150,18 @@ def set_current_catalog(self, catalog_name: str) -> None:
             https://learn.microsoft.com/en-us/fabric/data-warehouse/sql-query-editor#limitations
         """
         current_catalog = self.get_current_catalog()
+        target_catalog = self._normalize_catalog(catalog_name)
 
         # If already using the requested catalog, do nothing
-        if current_catalog and current_catalog == catalog_name:
-            logger.debug(f"Already using catalog '{catalog_name}', no action needed")
+        if current_catalog == target_catalog:
+            logger.debug("Already using the requested Fabric catalog state, no action needed")
             return
 
-        logger.info(f"Switching from catalog '{current_catalog}' to '{catalog_name}'")
+        logger.info(
+            "Switching from catalog '%s' to '%s'",
+            self._catalog_state_label(current_catalog),
+            self._catalog_state_label(target_catalog),
+        )
 
         # commit the transaction before closing the connection to help prevent errors like:
         # > Snapshot isolation transaction failed in database because the object accessed by the statement has been modified by a
@@ -149,14 +172,14 @@ def set_current_catalog(self, catalog_name: str) -> None:
         # note: we call close() on the connection pool instead of self.close() because self.close() calls close_all()
         # on the connection pool but we just want to close the connection for this thread
         self._connection_pool.close()
-        self._target_catalog = catalog_name  # new connections will use this catalog
+        self._target_catalog = target_catalog
 
         catalog_after_switch = self.get_current_catalog()
 
-        if catalog_after_switch != catalog_name:
+        if catalog_after_switch != target_catalog:
             # We need to raise an error if the catalog switch failed to prevent the operation that needed the catalog switch from being run against the wrong catalog
             raise SQLMeshError(
-                f"Unable to switch catalog to {catalog_name}, catalog ended up as {catalog_after_switch}"
+                f"Unable to switch catalog to {target_catalog}, catalog ended up as {catalog_after_switch}"
             )
 
     def alter_table(

tests/core/engine_adapter/test_fabric.py

@@ -19,15 +19,15 @@ def adapter(make_mocked_engine_adapter: t.Callable) -> FabricEngineAdapter:
     return make_mocked_engine_adapter(FabricEngineAdapter)
 
 
-def test_get_current_catalog_uses_target_catalog_or_configured_database(
+def test_get_current_catalog_uses_only_explicit_target_catalog(
     make_mocked_engine_adapter: t.Callable,
 ):
     adapter = make_mocked_engine_adapter(
         FabricEngineAdapter,
         database="default_catalog",
     )
 
-    assert adapter.get_current_catalog() == "default_catalog"
+    assert adapter.get_current_catalog() is None
 
     adapter._target_catalog = "switched_catalog"
 
@@ -36,7 +36,7 @@ def test_get_current_catalog_uses_target_catalog_or_configured_database(
     adapter._connection_pool.close()
 
     assert adapter._connection_pool.get_attribute("target_catalog") is None
-    assert adapter.get_current_catalog() == "default_catalog"
+    assert adapter.get_current_catalog() is None
     adapter.cursor.execute.assert_not_called()
 
 
@@ -63,6 +63,46 @@ def test_set_current_catalog_does_not_query_database(
     adapter.cursor.execute.assert_not_called()
 
 
+def test_set_current_catalog_to_default_clears_explicit_target(
+    make_mocked_engine_adapter: t.Callable,
+):
+    adapter = make_mocked_engine_adapter(
+        FabricEngineAdapter,
+        default_catalog="core",
+        database="core",
+    )
+
+    adapter.set_current_catalog("planning")
+    adapter.set_current_catalog("core")
+
+    assert adapter.get_current_catalog() is None
+    adapter.cursor.execute.assert_not_called()
+
+
+def test_catalog_scoped_call_restores_to_neutral_state(
+    make_mocked_engine_adapter: t.Callable,
+    mocker: MockerFixture,
+):
+    adapter = make_mocked_engine_adapter(
+        FabricEngineAdapter,
+        default_catalog="core",
+        database="core",
+    )
+    set_current_catalog_spy = mocker.patch.object(
+        adapter,
+        "set_current_catalog",
+        wraps=adapter.set_current_catalog,
+    )
+    adapter.cursor.fetchone.return_value = (1,)
+
+    adapter.table_exists("planning.db.table")
+
+    assert [
+        call.args[0] for call in set_current_catalog_spy.call_args_list
+    ] == ["planning", None]
+    assert adapter.get_current_catalog() is None
+
+
 def test_columns(adapter: FabricEngineAdapter):
     adapter.cursor.fetchall.return_value = [
         ("decimal_ps", "decimal", None, 5, 4),