fix(fabric): object list now contains catalog

Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com>

Author: cmgoffena13

sqlmesh/core/engine_adapter/fabric.py

@@ -8,6 +8,7 @@
 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,
@@ -85,6 +86,14 @@ def _catalog_state_label(self, catalog_name: t.Optional[str]) -> str:
             or "<default>"
         )
 
+    def _resolved_catalog(self) -> t.Optional[str]:
+        return (
+            self.get_current_catalog()
+            or self._normalize_catalog(self._connected_catalog)
+            or self._default_catalog
+            or self._extra_config.get("database")
+        )
+
     @property
     def api_client(self) -> FabricHttpClient:
         # the requests Session is not guaranteed to be threadsafe
@@ -224,18 +233,18 @@ 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
+    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]:
-        objects = super()._get_data_objects(schema_name, object_names)
         # Fabric uses None as "default catalog" so we skip reconnects. Other engines
-        # return a real warehouse name here; fill it in so this listing matches them
-        # and so the data-object cache key matches lookups that use that name.
-        catalog = (
-            self.get_current_catalog()
-            or self._default_catalog
-            or self._extra_config.get("database")
-        )
+        # return a real warehouse name here.
+        # Cache on schema.table due to @set_catalog stripping the catalog. Then put a
+        # warehouse name on the returned objects so listing matches other engines.
+        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 [

tests/core/engine_adapter/test_fabric.py

@@ -19,6 +19,18 @@ def adapter(make_mocked_engine_adapter: t.Callable) -> FabricEngineAdapter:
     return make_mocked_engine_adapter(FabricEngineAdapter)
 
 
+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_uses_only_explicit_target_catalog(
     make_mocked_engine_adapter: t.Callable,
 ):
@@ -471,11 +483,129 @@ def test_get_data_objects_uses_default_catalog_when_current_is_none(
         return_value=pd.DataFrame([{"name": "test_table", "schema_name": "dbo", "type": "TABLE"}]),
     )
 
-    assert adapter._get_data_objects("dbo") == [
+    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 catalog is None; the connection is still on planning.
+
+    Unqualified list must be planning. 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.get_current_catalog() is None
+
+    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)