Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions sqlmesh/core/engine_adapter/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
88 changes: 84 additions & 4 deletions sqlmesh/core/engine_adapter/fabric.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,18 @@
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
from sqlmesh.utils.connection_pool import ConnectionPool
from sqlmesh.core.schema_diff import TableAlterOperation
from sqlmesh.utils import random_id


logger = logging.getLogger(__name__)


Expand Down Expand Up @@ -85,6 +86,15 @@ def _catalog_state_label(self, catalog_name: t.Optional[str]) -> str:
or "<default>"
)

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
Expand Down Expand Up @@ -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.
Expand All @@ -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")
Expand Down Expand Up @@ -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]],
Expand Down
7 changes: 5 additions & 2 deletions sqlmesh/core/engine_adapter/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading
Loading