diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 930cdf7cd4..70eca19a71 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -457,6 +457,14 @@ def get_current_catalog(self) -> t.Optional[str]: """Returns the catalog name of the current connection.""" raise NotImplementedError() + def _get_current_catalog_for_set_catalog(self) -> t.Optional[str]: + """Catalog token @set_catalog saves and restores. + + Override when get_current_catalog() reports a default catalog name but + the restore must keep an internal sentinel so the connection is not reopened. + """ + return self.get_current_catalog() + def set_current_catalog(self, catalog: str) -> None: """Sets the catalog name of the current connection.""" raise NotImplementedError() diff --git a/sqlmesh/core/engine_adapter/fabric.py b/sqlmesh/core/engine_adapter/fabric.py index 97fe5b5008..4f687bbe9e 100644 --- a/sqlmesh/core/engine_adapter/fabric.py +++ b/sqlmesh/core/engine_adapter/fabric.py @@ -8,9 +8,11 @@ from sqlglot import exp from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_result from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter +from sqlmesh.core.dialect import to_schema from sqlmesh.core.engine_adapter.shared import ( CommentCreationTable, CommentCreationView, + DataObject, InsertOverwriteStrategy, ) from sqlmesh.utils.errors import SQLMeshError @@ -18,7 +20,6 @@ from sqlmesh.core.schema_diff import TableAlterOperation from sqlmesh.utils import random_id - logger = logging.getLogger(__name__) @@ -85,6 +86,15 @@ def _catalog_state_label(self, catalog_name: t.Optional[str]) -> str: or "" ) + def _resolved_catalog(self) -> t.Optional[str]: + # A lazy restore sets the logical target back to the default but leaves + # the open connection on the previous warehouse. Name that connection, + # because that is the warehouse the query hits. + connected_catalog = self._normalize_catalog(self._connected_catalog) + if self._normalize_catalog(self._target_catalog) is None and connected_catalog: + return connected_catalog + return self.get_current_catalog() + @property def api_client(self) -> FabricHttpClient: # the requests Session is not guaranteed to be threadsafe @@ -145,10 +155,29 @@ def _drop_catalog(self, catalog_name: exp.Identifier) -> None: if warehouse_name in in_use: self.close() - def get_current_catalog(self) -> t.Optional[str]: - """Return the explicit Fabric catalog target for the current thread.""" + def _get_current_catalog_for_set_catalog(self) -> t.Optional[str]: + """Override of EngineAdapter._get_current_catalog_for_set_catalog. + + The base method returns get_current_catalog(). @set_catalog saves that + value and writes it back with set_current_catalog. Fabric's public method + returns the configured warehouse name, so restoring it would reconnect. + This returns the normalized internal target instead. None is that default, + and restoring None skips the reconnect. + """ return self._normalize_catalog(self._target_catalog) + def get_current_catalog(self) -> t.Optional[str]: + """Return the warehouse name for the current thread. + + The configured default stays None internally so set_current_catalog can + skip a reconnect. Callers still receive that warehouse's name. None is + returned only when no warehouse is configured. + """ + explicit_catalog = self._normalize_catalog(self._target_catalog) + if explicit_catalog: + return explicit_catalog + return self._default_catalog or self._extra_config.get("database") + def set_current_catalog(self, catalog_name: t.Optional[str]) -> None: """ Set the current catalog for Microsoft Fabric connections. @@ -173,10 +202,12 @@ def set_current_catalog(self, catalog_name: t.Optional[str]) -> None: explicit_default_catalog = catalog_name is not None and target_catalog is None connected_catalog = self._normalize_catalog(self._connected_catalog) + # Compare the internal sentinel, not get_current_catalog(). That method + # returns the configured warehouse name when this sentinel is None. # An explicit request for the default catalog must also match the catalog # used by the open connection. A lazy restore with None only updates the # logical target and intentionally leaves that connection in place. - if self.get_current_catalog() == target_catalog and ( + if self._normalize_catalog(self._target_catalog) == target_catalog and ( not explicit_default_catalog or connected_catalog is None ): logger.debug("Already using requested Fabric catalog state, no action needed") @@ -224,6 +255,55 @@ def set_current_catalog(self, catalog_name: t.Optional[str]) -> None: self._target_catalog = target_catalog + def _get_data_objects( + self, schema_name: t.Union[str, exp.Table], object_names: t.Optional[t.Set[str]] = None + ) -> t.List[DataObject]: + objects = super()._get_data_objects(schema_name, object_names) + # Runs before EngineAdapter.get_data_objects writes the cache. @set_catalog + # has already stripped the catalog, so the lookup key is schema.table. + # MSSQL stamps get_current_catalog() here, which is now the default warehouse + # name, and that would store catalog.schema.table instead. Blank that name so the + # stored key matches. get_data_objects puts it back on the objects it returns. + default_catalog = self._default_catalog or self._extra_config.get("database") + if not default_catalog: + return objects + return [ + DataObject( + catalog=None if obj.catalog == default_catalog else obj.catalog, + schema=obj.schema_name, + name=obj.name, + type=obj.type, + clustering_key=obj.clustering_key, + ) + for obj in objects + ] + + def get_data_objects( + self, + schema_name: t.Union[str, exp.Table], + object_names: t.Optional[t.Set[str]] = None, + safe_to_cache: bool = False, + ) -> t.List[DataObject]: + # Runs after the cache read. _get_data_objects blanked the default warehouse + # so the cached key stays schema.table. Callers still need that name on the + # returned objects. A schema that already names a warehouse wins; otherwise + # use the connection actually queried, which may still be the previous + # warehouse after a lazy restore. + objects = super().get_data_objects(schema_name, object_names, safe_to_cache) + catalog = to_schema(schema_name).catalog or self._resolved_catalog() + if not catalog: + return objects + return [ + DataObject( + catalog=obj.catalog or catalog, + schema=obj.schema_name, + name=obj.name, + type=obj.type, + clustering_key=obj.clustering_key, + ) + for obj in objects + ] + def alter_table( self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]], diff --git a/sqlmesh/core/engine_adapter/shared.py b/sqlmesh/core/engine_adapter/shared.py index ba0e1fa619..04c56dae0b 100644 --- a/sqlmesh/core/engine_adapter/shared.py +++ b/sqlmesh/core/engine_adapter/shared.py @@ -348,8 +348,11 @@ def internal_wrapper(*args: t.Any, **kwargs: t.Any) -> t.Any: f"{engine_adapter.dialect} requires that all catalog operations be against a single catalog: {engine_adapter._default_catalog}. Provided catalog: {catalog_name}" ) return func(*list_args, **kwargs) - # Set the catalog name on the engine adapter if needed - current_catalog = engine_adapter.get_current_catalog() + # Set the catalog name on the engine adapter if needed. + # See EngineAdapter._get_current_catalog_for_set_catalog. Fabric reports + # the configured warehouse name from get_current_catalog() and keeps None + # internally, so a restore of the default warehouse can skip a reconnect. + current_catalog = engine_adapter._get_current_catalog_for_set_catalog() if catalog_name != current_catalog: engine_adapter.set_current_catalog(catalog_name) resp = func(*list_args, **kwargs) diff --git a/tests/core/engine_adapter/test_fabric.py b/tests/core/engine_adapter/test_fabric.py index a6760e9531..ceab0e7868 100644 --- a/tests/core/engine_adapter/test_fabric.py +++ b/tests/core/engine_adapter/test_fabric.py @@ -11,7 +11,7 @@ from sqlmesh.core.engine_adapter import FabricEngineAdapter from sqlmesh.core.engine_adapter.fabric import FabricHttpClient from tests.core.engine_adapter import to_sql_calls -from sqlmesh.core.engine_adapter.shared import DataObject +from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType pytestmark = [pytest.mark.engine, pytest.mark.fabric] @@ -41,7 +41,19 @@ def _conflict_response(mocker: MockerFixture, error_code: str) -> t.Any: return resp -def test_get_current_catalog_uses_only_explicit_target_catalog( +def _record_catalogs_at_execute(adapter: FabricEngineAdapter) -> t.List[t.Optional[str]]: + catalogs: t.List[t.Optional[str]] = [] + real_execute = adapter.execute + + def execute(*args: t.Any, **kwargs: t.Any) -> None: + catalogs.append(adapter._resolved_catalog()) + return real_execute(*args, **kwargs) + + adapter.execute = execute # type: ignore[method-assign] + return catalogs + + +def test_get_current_catalog_returns_configured_database_when_target_is_default( make_mocked_engine_adapter: t.Callable, ): adapter = make_mocked_engine_adapter( @@ -49,7 +61,10 @@ def test_get_current_catalog_uses_only_explicit_target_catalog( database="default_catalog", ) - assert adapter.get_current_catalog() is None + current_catalog = adapter.get_current_catalog() + assert current_catalog is not None + assert current_catalog.lower() == "default_catalog" + assert adapter._target_catalog is None adapter._target_catalog = "switched_catalog" @@ -58,10 +73,31 @@ def test_get_current_catalog_uses_only_explicit_target_catalog( adapter._connection_pool.close() assert adapter._connection_pool.get_attribute("target_catalog") is None - assert adapter.get_current_catalog() is None + assert adapter.get_current_catalog() == "default_catalog" adapter.cursor.execute.assert_not_called() +def test_shared_catalog_operations_can_lowercase_the_default_warehouse( + make_mocked_engine_adapter: t.Callable, +): + """test_catalog_operations lowercases the current catalog, then restores it.""" + adapter = make_mocked_engine_adapter( + FabricEngineAdapter, + database="ci_abc", + ) + + current_catalog = adapter.get_current_catalog().lower() + assert current_catalog == "ci_abc" + assert adapter._target_catalog is None + + adapter.set_current_catalog("testing") + assert adapter.get_current_catalog().lower() == "testing" + + adapter.set_current_catalog(current_catalog) + assert adapter.get_current_catalog().lower() == current_catalog + assert adapter._target_catalog is None + + def test_get_current_catalog_returns_none_without_target_or_database( make_mocked_engine_adapter: t.Callable, ): @@ -97,7 +133,8 @@ def test_set_current_catalog_to_default_clears_explicit_target( adapter.set_current_catalog("planning") adapter.set_current_catalog("core") - assert adapter.get_current_catalog() is None + assert adapter._target_catalog is None + assert adapter.get_current_catalog() == "core" adapter.cursor.execute.assert_not_called() @@ -119,8 +156,9 @@ def test_catalog_scoped_call_restores_to_neutral_without_close( # Decorator calls set_current_catalog("planning") then set_current_catalog(None). # Only the first call (None→planning) should trigger a connection close. assert close_spy.call_count == 1 + assert adapter._target_catalog is None assert adapter._connected_catalog == "planning" - assert adapter.get_current_catalog() is None + assert adapter.get_current_catalog() == "core" def test_default_catalog_after_non_default_catalog_reconnects( @@ -139,8 +177,9 @@ def test_default_catalog_after_non_default_catalog_reconnects( adapter.table_exists("core.db.table") assert close_spy.call_count == 2 + assert adapter._target_catalog is None assert adapter._connected_catalog is None - assert adapter.get_current_catalog() is None + assert adapter.get_current_catalog() == "core" def test_repeated_same_catalog_reuses_connection( @@ -494,3 +533,154 @@ def test_create_warehouse_unrelated_conflict_still_raises( with pytest.raises(requests.HTTPError): fabric_http_client.create_warehouse("my_warehouse") + + +def test_get_data_objects_uses_configured_default_catalog( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, +) -> None: + adapter = make_mocked_engine_adapter( + FabricEngineAdapter, + default_catalog="ci_abc", + database="ci_abc", + patch_get_data_objects=False, + ) + assert adapter._target_catalog is None + assert adapter.get_current_catalog() == "ci_abc" + + mocker.patch.object( + adapter, + "fetchdf", + return_value=pd.DataFrame([{"name": "test_table", "schema_name": "dbo", "type": "TABLE"}]), + ) + + assert adapter.get_data_objects("ci_abc.dbo") == [ + DataObject( + catalog="ci_abc", + schema="dbo", + name="test_table", + type=DataObjectType.TABLE, + ) + ] + + +def test_get_data_objects_cache_hits_for_default_catalog( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, +) -> None: + """Listing fills catalog on the return value. Cache must still hit.""" + adapter = make_mocked_engine_adapter( + FabricEngineAdapter, + default_catalog="ci_abc", + database="ci_abc", + patch_get_data_objects=False, + ) + fetchdf = mocker.patch.object( + adapter, + "fetchdf", + return_value=pd.DataFrame([{"name": "t", "schema_name": "dbo", "type": "TABLE"}]), + ) + + first = adapter.get_data_objects("ci_abc.dbo", {"t"}, safe_to_cache=True) + second = adapter.get_data_objects("ci_abc.dbo", {"t"}, safe_to_cache=True) + + assert first[0].catalog == "ci_abc" + assert second[0].catalog == "ci_abc" + assert fetchdf.call_count == 1 + + +def test_get_data_objects_labels_connected_warehouse_after_lazy_restore( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, +) -> None: + """Logical target is the default; the connection is still on planning. + + The public catalog is the configured warehouse. The unqualified list must + still be planning, because that is the warehouse the query hits. + ci_abc.dbo must still switch. + """ + adapter = make_mocked_engine_adapter( + FabricEngineAdapter, + default_catalog="ci_abc", + database="ci_abc", + patch_get_data_objects=False, + ) + adapter.set_current_catalog("planning") + adapter.set_current_catalog(None) + assert adapter._target_catalog is None + assert adapter._connected_catalog == "planning" + assert adapter.get_current_catalog() == "ci_abc" + + fetchdf = mocker.patch.object( + adapter, + "fetchdf", + return_value=pd.DataFrame([{"name": "t", "schema_name": "dbo", "type": "TABLE"}]), + ) + + objects = adapter.get_data_objects("dbo") + + assert objects == [ + DataObject( + catalog="planning", + schema="dbo", + name="t", + type=DataObjectType.TABLE, + ) + ] + + assert adapter.get_data_objects("ci_abc.dbo") == [ + DataObject( + catalog="ci_abc", + schema="dbo", + name="t", + type=DataObjectType.TABLE, + ) + ] + assert fetchdf.call_count == 2 + + +def test_drop_data_object_default_catalog_drops_without_requalifying( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, +) -> None: + """DROP SQL has no warehouse. Prove it ran on ci_abc, without reconnecting.""" + adapter = make_mocked_engine_adapter( + FabricEngineAdapter, + default_catalog="ci_abc", + database="ci_abc", + ) + close = mocker.spy(adapter._connection_pool, "close") + catalogs_at_execute = _record_catalogs_at_execute(adapter) + + adapter.drop_data_object( + DataObject(catalog="ci_abc", schema="dbo", name="v", type=DataObjectType.VIEW) + ) + + close.assert_not_called() + assert catalogs_at_execute == ["ci_abc"] + assert to_sql_calls(adapter) == ["DROP VIEW IF EXISTS [dbo].[v];"] + + +def test_drop_data_object_default_catalog_reconnects_when_connected_elsewhere( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, +) -> None: + """DROP SQL has no warehouse. Prove it ran on ci_abc, then restored planning.""" + adapter = make_mocked_engine_adapter( + FabricEngineAdapter, + default_catalog="ci_abc", + database="ci_abc", + ) + adapter.set_current_catalog("planning") + assert adapter.get_current_catalog() == "planning" + close = mocker.spy(adapter._connection_pool, "close") + catalogs_at_execute = _record_catalogs_at_execute(adapter) + + adapter.drop_data_object( + DataObject(catalog="ci_abc", schema="dbo", name="v", type=DataObjectType.VIEW) + ) + + assert catalogs_at_execute == ["ci_abc"] + assert adapter.get_current_catalog() == "planning" + assert close.call_count == 2 + assert "DROP VIEW IF EXISTS [dbo].[v];" in to_sql_calls(adapter)