Skip to content

Commit c491e42

Browse files
committed
fix(fabric): align fabric adapter output to adapter interface
Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com>
1 parent a45e6d6 commit c491e42

4 files changed

Lines changed: 116 additions & 26 deletions

File tree

‎sqlmesh/core/engine_adapter/base.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,14 @@ def get_current_catalog(self) -> t.Optional[str]:
457457
"""Returns the catalog name of the current connection."""
458458
raise NotImplementedError()
459459

460+
def _get_current_catalog_for_set_catalog(self) -> t.Optional[str]:
461+
"""Catalog token @set_catalog saves and restores.
462+
463+
Override when get_current_catalog() reports a default catalog name but
464+
the restore must keep an internal sentinel so the connection is not reopened.
465+
"""
466+
return self.get_current_catalog()
467+
460468
def set_current_catalog(self, catalog: str) -> None:
461469
"""Sets the catalog name of the current connection."""
462470
raise NotImplementedError()

‎sqlmesh/core/engine_adapter/fabric.py‎

Lines changed: 59 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -87,12 +87,13 @@ def _catalog_state_label(self, catalog_name: t.Optional[str]) -> str:
8787
)
8888

8989
def _resolved_catalog(self) -> t.Optional[str]:
90-
return (
91-
self.get_current_catalog()
92-
or self._normalize_catalog(self._connected_catalog)
93-
or self._default_catalog
94-
or self._extra_config.get("database")
95-
)
90+
# A lazy restore sets the logical target back to the default but leaves
91+
# the open connection on the previous warehouse. Name that connection,
92+
# because that is the warehouse the query hits.
93+
connected_catalog = self._normalize_catalog(self._connected_catalog)
94+
if self._normalize_catalog(self._target_catalog) is None and connected_catalog:
95+
return connected_catalog
96+
return self.get_current_catalog()
9697

9798
@property
9899
def api_client(self) -> FabricHttpClient:
@@ -154,10 +155,29 @@ def _drop_catalog(self, catalog_name: exp.Identifier) -> None:
154155
if warehouse_name in in_use:
155156
self.close()
156157

157-
def get_current_catalog(self) -> t.Optional[str]:
158-
"""Return the explicit Fabric catalog target for the current thread."""
158+
def _get_current_catalog_for_set_catalog(self) -> t.Optional[str]:
159+
"""Override of EngineAdapter._get_current_catalog_for_set_catalog.
160+
161+
The base method returns get_current_catalog(). @set_catalog saves that
162+
value and writes it back with set_current_catalog. Fabric's public method
163+
returns the configured warehouse name, so restoring it would reconnect.
164+
This returns the normalized internal target instead. None is that default,
165+
and restoring None skips the reconnect.
166+
"""
159167
return self._normalize_catalog(self._target_catalog)
160168

169+
def get_current_catalog(self) -> t.Optional[str]:
170+
"""Return the warehouse name for the current thread.
171+
172+
The configured default stays None internally so set_current_catalog can
173+
skip a reconnect. Callers still receive that warehouse's name. None is
174+
returned only when no warehouse is configured.
175+
"""
176+
explicit_catalog = self._normalize_catalog(self._target_catalog)
177+
if explicit_catalog:
178+
return explicit_catalog
179+
return self._default_catalog or self._extra_config.get("database")
180+
161181
def set_current_catalog(self, catalog_name: t.Optional[str]) -> None:
162182
"""
163183
Set the current catalog for Microsoft Fabric connections.
@@ -182,10 +202,12 @@ def set_current_catalog(self, catalog_name: t.Optional[str]) -> None:
182202
explicit_default_catalog = catalog_name is not None and target_catalog is None
183203
connected_catalog = self._normalize_catalog(self._connected_catalog)
184204

205+
# Compare the internal sentinel, not get_current_catalog(). That method
206+
# returns the configured warehouse name when this sentinel is None.
185207
# An explicit request for the default catalog must also match the catalog
186208
# used by the open connection. A lazy restore with None only updates the
187209
# logical target and intentionally leaves that connection in place.
188-
if self.get_current_catalog() == target_catalog and (
210+
if self._normalize_catalog(self._target_catalog) == target_catalog and (
189211
not explicit_default_catalog or connected_catalog is None
190212
):
191213
logger.debug("Already using requested Fabric catalog state, no action needed")
@@ -233,16 +255,40 @@ def set_current_catalog(self, catalog_name: t.Optional[str]) -> None:
233255

234256
self._target_catalog = target_catalog
235257

258+
def _get_data_objects(
259+
self, schema_name: t.Union[str, exp.Table], object_names: t.Optional[t.Set[str]] = None
260+
) -> t.List[DataObject]:
261+
objects = super()._get_data_objects(schema_name, object_names)
262+
# Runs before EngineAdapter.get_data_objects writes the cache. @set_catalog
263+
# has already stripped the catalog, so the lookup key is schema.table.
264+
# MSSQL stamps get_current_catalog() here, which is now the default warehouse
265+
# name, and that would store catalog.schema.table instead. Blank that name so the
266+
# stored key matches. get_data_objects puts it back on the objects it returns.
267+
default_catalog = self._default_catalog or self._extra_config.get("database")
268+
if not default_catalog:
269+
return objects
270+
return [
271+
DataObject(
272+
catalog=None if obj.catalog == default_catalog else obj.catalog,
273+
schema=obj.schema_name,
274+
name=obj.name,
275+
type=obj.type,
276+
clustering_key=obj.clustering_key,
277+
)
278+
for obj in objects
279+
]
280+
236281
def get_data_objects(
237282
self,
238283
schema_name: t.Union[str, exp.Table],
239284
object_names: t.Optional[t.Set[str]] = None,
240285
safe_to_cache: bool = False,
241286
) -> t.List[DataObject]:
242-
# Fabric uses None as "default catalog" so we skip reconnects. Other engines
243-
# return a real warehouse name here.
244-
# Cache on schema.table due to @set_catalog stripping the catalog. Then put a
245-
# warehouse name on the returned objects so listing matches other engines.
287+
# Runs after the cache read. _get_data_objects blanked the default warehouse
288+
# so the cached key stays schema.table. Callers still need that name on the
289+
# returned objects. A schema that already names a warehouse wins; otherwise
290+
# use the connection actually queried, which may still be the previous
291+
# warehouse after a lazy restore.
246292
objects = super().get_data_objects(schema_name, object_names, safe_to_cache)
247293
catalog = to_schema(schema_name).catalog or self._resolved_catalog()
248294
if not catalog:

‎sqlmesh/core/engine_adapter/shared.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -348,8 +348,11 @@ def internal_wrapper(*args: t.Any, **kwargs: t.Any) -> t.Any:
348348
f"{engine_adapter.dialect} requires that all catalog operations be against a single catalog: {engine_adapter._default_catalog}. Provided catalog: {catalog_name}"
349349
)
350350
return func(*list_args, **kwargs)
351-
# Set the catalog name on the engine adapter if needed
352-
current_catalog = engine_adapter.get_current_catalog()
351+
# Set the catalog name on the engine adapter if needed.
352+
# See EngineAdapter._get_current_catalog_for_set_catalog. Fabric reports
353+
# the configured warehouse name from get_current_catalog() and keeps None
354+
# internally, so a restore of the default warehouse can skip a reconnect.
355+
current_catalog = engine_adapter._get_current_catalog_for_set_catalog()
353356
if catalog_name != current_catalog:
354357
engine_adapter.set_current_catalog(catalog_name)
355358
resp = func(*list_args, **kwargs)

‎tests/core/engine_adapter/test_fabric.py‎

Lines changed: 44 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ def _conflict_response(mocker: MockerFixture, error_code: str) -> t.Any:
4040
resp.raise_for_status.side_effect = requests.HTTPError(response=resp)
4141
return resp
4242

43+
4344
def _record_catalogs_at_execute(adapter: FabricEngineAdapter) -> t.List[t.Optional[str]]:
4445
catalogs: t.List[t.Optional[str]] = []
4546
real_execute = adapter.execute
@@ -52,15 +53,18 @@ def execute(*args: t.Any, **kwargs: t.Any) -> None:
5253
return catalogs
5354

5455

55-
def test_get_current_catalog_uses_only_explicit_target_catalog(
56+
def test_get_current_catalog_returns_configured_database_when_target_is_default(
5657
make_mocked_engine_adapter: t.Callable,
5758
):
5859
adapter = make_mocked_engine_adapter(
5960
FabricEngineAdapter,
6061
database="default_catalog",
6162
)
6263

63-
assert adapter.get_current_catalog() is None
64+
current_catalog = adapter.get_current_catalog()
65+
assert current_catalog is not None
66+
assert current_catalog.lower() == "default_catalog"
67+
assert adapter._target_catalog is None
6468

6569
adapter._target_catalog = "switched_catalog"
6670

@@ -69,10 +73,31 @@ def test_get_current_catalog_uses_only_explicit_target_catalog(
6973
adapter._connection_pool.close()
7074

7175
assert adapter._connection_pool.get_attribute("target_catalog") is None
72-
assert adapter.get_current_catalog() is None
76+
assert adapter.get_current_catalog() == "default_catalog"
7377
adapter.cursor.execute.assert_not_called()
7478

7579

80+
def test_shared_catalog_operations_can_lowercase_the_default_warehouse(
81+
make_mocked_engine_adapter: t.Callable,
82+
):
83+
"""test_catalog_operations lowercases the current catalog, then restores it."""
84+
adapter = make_mocked_engine_adapter(
85+
FabricEngineAdapter,
86+
database="ci_abc",
87+
)
88+
89+
current_catalog = adapter.get_current_catalog().lower()
90+
assert current_catalog == "ci_abc"
91+
assert adapter._target_catalog is None
92+
93+
adapter.set_current_catalog("testing")
94+
assert adapter.get_current_catalog().lower() == "testing"
95+
96+
adapter.set_current_catalog(current_catalog)
97+
assert adapter.get_current_catalog().lower() == current_catalog
98+
assert adapter._target_catalog is None
99+
100+
76101
def test_get_current_catalog_returns_none_without_target_or_database(
77102
make_mocked_engine_adapter: t.Callable,
78103
):
@@ -108,7 +133,8 @@ def test_set_current_catalog_to_default_clears_explicit_target(
108133
adapter.set_current_catalog("planning")
109134
adapter.set_current_catalog("core")
110135

111-
assert adapter.get_current_catalog() is None
136+
assert adapter._target_catalog is None
137+
assert adapter.get_current_catalog() == "core"
112138
adapter.cursor.execute.assert_not_called()
113139

114140

@@ -130,8 +156,9 @@ def test_catalog_scoped_call_restores_to_neutral_without_close(
130156
# Decorator calls set_current_catalog("planning") then set_current_catalog(None).
131157
# Only the first call (None→planning) should trigger a connection close.
132158
assert close_spy.call_count == 1
159+
assert adapter._target_catalog is None
133160
assert adapter._connected_catalog == "planning"
134-
assert adapter.get_current_catalog() is None
161+
assert adapter.get_current_catalog() == "core"
135162

136163

137164
def test_default_catalog_after_non_default_catalog_reconnects(
@@ -150,8 +177,9 @@ def test_default_catalog_after_non_default_catalog_reconnects(
150177
adapter.table_exists("core.db.table")
151178

152179
assert close_spy.call_count == 2
180+
assert adapter._target_catalog is None
153181
assert adapter._connected_catalog is None
154-
assert adapter.get_current_catalog() is None
182+
assert adapter.get_current_catalog() == "core"
155183

156184

157185
def test_repeated_same_catalog_reuses_connection(
@@ -507,7 +535,7 @@ def test_create_warehouse_unrelated_conflict_still_raises(
507535
fabric_http_client.create_warehouse("my_warehouse")
508536

509537

510-
def test_get_data_objects_uses_default_catalog_when_current_is_none(
538+
def test_get_data_objects_uses_configured_default_catalog(
511539
make_mocked_engine_adapter: t.Callable,
512540
mocker: MockerFixture,
513541
) -> None:
@@ -517,7 +545,8 @@ def test_get_data_objects_uses_default_catalog_when_current_is_none(
517545
database="ci_abc",
518546
patch_get_data_objects=False,
519547
)
520-
assert adapter.get_current_catalog() is None
548+
assert adapter._target_catalog is None
549+
assert adapter.get_current_catalog() == "ci_abc"
521550

522551
mocker.patch.object(
523552
adapter,
@@ -564,9 +593,11 @@ def test_get_data_objects_labels_connected_warehouse_after_lazy_restore(
564593
make_mocked_engine_adapter: t.Callable,
565594
mocker: MockerFixture,
566595
) -> None:
567-
"""Logical catalog is None; the connection is still on planning.
596+
"""Logical target is the default; the connection is still on planning.
568597
569-
Unqualified list must be planning. ci_abc.dbo must still switch.
598+
The public catalog is the configured warehouse. The unqualified list must
599+
still be planning, because that is the warehouse the query hits.
600+
ci_abc.dbo must still switch.
570601
"""
571602
adapter = make_mocked_engine_adapter(
572603
FabricEngineAdapter,
@@ -576,7 +607,9 @@ def test_get_data_objects_labels_connected_warehouse_after_lazy_restore(
576607
)
577608
adapter.set_current_catalog("planning")
578609
adapter.set_current_catalog(None)
579-
assert adapter.get_current_catalog() is None
610+
assert adapter._target_catalog is None
611+
assert adapter._connected_catalog == "planning"
612+
assert adapter.get_current_catalog() == "ci_abc"
580613

581614
fetchdf = mocker.patch.object(
582615
adapter,

0 commit comments

Comments
 (0)