sqlmesh.core.engine_adapter.fabric

  1from __future__ import annotations
  2
  3import typing as t
  4import logging
  5import requests
  6import time
  7from functools import cached_property
  8from sqlglot import exp
  9from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_result
 10from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
 11from sqlmesh.core.engine_adapter.shared import (
 12    CommentCreationTable,
 13    CommentCreationView,
 14    InsertOverwriteStrategy,
 15)
 16from sqlmesh.utils.errors import SQLMeshError
 17from sqlmesh.utils.connection_pool import ConnectionPool
 18from sqlmesh.core.schema_diff import TableAlterOperation
 19from sqlmesh.utils import random_id
 20
 21
 22logger = logging.getLogger(__name__)
 23
 24
 25class FabricEngineAdapter(MSSQLEngineAdapter):
 26    """
 27    Adapter for Microsoft Fabric.
 28    """
 29
 30    DIALECT = "fabric"
 31    SUPPORTS_INDEXES = False
 32    SUPPORTS_TRANSACTIONS = False
 33    SUPPORTS_CREATE_DROP_CATALOG = True
 34    INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT
 35    # There is no standard method to handle comments in Fabric for now, so we disable it.
 36    # Otherwise, it would be inherited from MSSQL and would not work.
 37    COMMENT_CREATION_TABLE = CommentCreationTable.UNSUPPORTED
 38    COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
 39
 40    def __init__(
 41        self, connection_factory_or_pool: t.Union[t.Callable, t.Any], *args: t.Any, **kwargs: t.Any
 42    ) -> None:
 43        # Wrap connection factory to support changing the catalog dynamically at runtime
 44        if not isinstance(connection_factory_or_pool, ConnectionPool):
 45            original_connection_factory = connection_factory_or_pool
 46
 47            connection_factory_or_pool = lambda *args, **kwargs: original_connection_factory(
 48                target_catalog=self._target_catalog, *args, **kwargs
 49            )
 50
 51        super().__init__(connection_factory_or_pool, *args, **kwargs)
 52
 53    @property
 54    def _target_catalog(self) -> t.Optional[str]:
 55        return self._connection_pool.get_attribute("target_catalog")
 56
 57    @_target_catalog.setter
 58    def _target_catalog(self, value: t.Optional[str]) -> None:
 59        self._connection_pool.set_attribute("target_catalog", value)
 60
 61    @property
 62    def _connected_catalog(self) -> t.Optional[str]:
 63        """Catalog the currently-open thread-local connection is actually using."""
 64        return self._connection_pool.get_attribute("connected_catalog")
 65
 66    @_connected_catalog.setter
 67    def _connected_catalog(self, value: t.Optional[str]) -> None:
 68        self._connection_pool.set_attribute("connected_catalog", value)
 69
 70    def _normalize_catalog(self, catalog_name: t.Optional[str]) -> t.Optional[str]:
 71        if not catalog_name:
 72            return None
 73
 74        default_catalog = self._default_catalog or self._extra_config.get("database")
 75        if default_catalog and catalog_name == default_catalog:
 76            return None
 77
 78        return catalog_name
 79
 80    def _catalog_state_label(self, catalog_name: t.Optional[str]) -> str:
 81        return (
 82            catalog_name
 83            or self._default_catalog
 84            or self._extra_config.get("database")
 85            or "<default>"
 86        )
 87
 88    @property
 89    def api_client(self) -> FabricHttpClient:
 90        # the requests Session is not guaranteed to be threadsafe
 91        # so we create a http client per thread on demand
 92        if existing_client := self._connection_pool.get_attribute("api_client"):
 93            return existing_client
 94
 95        tenant_id: t.Optional[str] = self._extra_config.get("tenant_id")
 96        workspace_id: t.Optional[str] = self._extra_config.get("workspace_id")
 97        client_id: t.Optional[str] = self._extra_config.get("user")
 98        client_secret: t.Optional[str] = self._extra_config.get("password")
 99
100        if not tenant_id or not client_id or not client_secret:
101            raise SQLMeshError(
102                "Service Principal authentication requires tenant_id, client_id, and client_secret "
103                "in the Fabric connection configuration"
104            )
105
106        if not workspace_id:
107            raise SQLMeshError(
108                "Fabric requires the workspace_id to be configured in the connection configuration to create / drop catalogs"
109            )
110
111        client = FabricHttpClient(
112            tenant_id=tenant_id,
113            workspace_id=workspace_id,
114            client_id=client_id,
115            client_secret=client_secret,
116        )
117
118        self._connection_pool.set_attribute("api_client", client)
119        return client
120
121    def _create_catalog(self, catalog_name: exp.Identifier) -> None:
122        """Create a catalog (warehouse) in Microsoft Fabric via REST API."""
123        warehouse_name = catalog_name.sql(dialect=self.dialect, identify=False)
124        logger.info(f"Creating Fabric warehouse: {warehouse_name}")
125
126        self.api_client.create_warehouse(warehouse_name)
127
128    def _drop_catalog(self, catalog_name: exp.Identifier) -> None:
129        """Drop a catalog (warehouse) in Microsoft Fabric via REST API."""
130        warehouse_name = catalog_name.sql(dialect=self.dialect, identify=False)
131
132        logger.info(f"Deleting Fabric warehouse: {warehouse_name}")
133        self.api_client.delete_warehouse(warehouse_name)
134
135        # Close all connections if any thread may be using the dropped warehouse.
136        # We must check both the logical target and the physical connection catalog
137        # (falling back to the configured default when either is neutral) because
138        # Fabric validates the DATABASE= connection argument and raises
139        # 'Authentication Failed' when it points at a non-existent warehouse.
140        default_db = self._extra_config.get("database")
141        in_use = {
142            self.get_current_catalog() or default_db,
143            self._normalize_catalog(self._connected_catalog) or default_db,
144        }
145        if warehouse_name in in_use:
146            self.close()
147
148    def get_current_catalog(self) -> t.Optional[str]:
149        """Return the explicit Fabric catalog target for the current thread."""
150        return self._normalize_catalog(self._target_catalog)
151
152    def set_current_catalog(self, catalog_name: t.Optional[str]) -> None:
153        """
154        Set the current catalog for Microsoft Fabric connections.
155
156        Override to handle Fabric's stateless session limitation where USE statements
157        don't persist across queries. Instead, we close existing connections and
158        recreate them with the new catalog in the connection configuration.
159
160        Args:
161            catalog_name: The name of the catalog (warehouse) to switch to.
162                The configured default catalog is treated as the neutral state.
163
164        Note:
165            Fabric doesn't support catalog switching via USE statements because each
166            statement runs as an independent session. This method works around this
167            limitation by updating the connection pool with new catalog configuration.
168
169        See:
170            https://learn.microsoft.com/en-us/fabric/data-warehouse/sql-query-editor#limitations
171        """
172        target_catalog = self._normalize_catalog(catalog_name)
173        explicit_default_catalog = catalog_name is not None and target_catalog is None
174        connected_catalog = self._normalize_catalog(self._connected_catalog)
175
176        # An explicit request for the default catalog must also match the catalog
177        # used by the open connection. A lazy restore with None only updates the
178        # logical target and intentionally leaves that connection in place.
179        if self.get_current_catalog() == target_catalog and (
180            not explicit_default_catalog or connected_catalog is None
181        ):
182            logger.debug("Already using requested Fabric catalog state, no action needed")
183            return
184
185        # Decide whether the open connection needs to be replaced.
186        #
187        # The set_catalog decorator restores the previous catalog (often None)
188        # after every catalog-scoped call.  For Fabric, a connection close +
189        # reopen is expensive because each new connection goes through ODBC and
190        # the Fabric gateway.  We therefore apply lazy connection management:
191        #
192        #  * When restoring to neutral (target=None): just update _target_catalog.
193        #    The existing connection stays alive and will be reused or replaced
194        #    on the next real switch, avoiding a pointless bounce through the
195        #    default catalog.
196        #
197        #  * When switching to a non-neutral catalog: only close/reopen if the
198        #    open connection is already on a different catalog.  If a previous
199        #    restore-to-neutral left the connection on the right catalog, we
200        #    skip the close entirely.
201        needs_reconnect = (target_catalog is not None or explicit_default_catalog) and (
202            connected_catalog != target_catalog
203        )
204
205        if needs_reconnect:
206            logger.info(
207                "Switching connection from catalog '%s' to '%s'",
208                self._catalog_state_label(connected_catalog),
209                self._catalog_state_label(target_catalog),
210            )
211            # Commit before closing to avoid snapshot-isolation errors on
212            # subsequent queries in the new connection.
213            self._connection_pool.commit()
214            # note: close() on the pool (not self.close()) to only affect this
215            # thread's connection rather than all threads.
216            self._connection_pool.close()
217            self._connected_catalog = target_catalog
218        else:
219            logger.debug(
220                "Updating catalog target to '%s' (connection remains on '%s')",
221                self._catalog_state_label(target_catalog),
222                self._catalog_state_label(connected_catalog),
223            )
224
225        self._target_catalog = target_catalog
226
227    def alter_table(
228        self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]]
229    ) -> None:
230        """
231        Applies alter expressions to a table. Fabric has limited support for ALTER TABLE,
232        so this method implements a workaround for column type changes.
233        This method is self-contained and sets its own catalog context.
234        """
235        if not alter_expressions:
236            return
237
238        # Get the target table from the first expression to determine the correct catalog.
239        first_op = alter_expressions[0]
240        expression = first_op.expression if isinstance(first_op, TableAlterOperation) else first_op
241        if not isinstance(expression, exp.Alter) or not expression.this.catalog:
242            # Fallback for unexpected scenarios
243            logger.warning(
244                "Could not determine catalog from alter expression, executing with current context."
245            )
246            super().alter_table(alter_expressions)
247            return
248
249        target_catalog = expression.this.catalog
250        self.set_current_catalog(target_catalog)
251
252        with self.transaction():
253            for op in alter_expressions:
254                expression = op.expression if isinstance(op, TableAlterOperation) else op
255
256                if not isinstance(expression, exp.Alter):
257                    self.execute(expression)
258                    continue
259
260                for action in expression.actions:
261                    table_name = expression.this
262
263                    table_name_without_catalog = table_name.copy()
264                    table_name_without_catalog.set("catalog", None)
265
266                    is_type_change = isinstance(action, exp.AlterColumn) and action.args.get(
267                        "dtype"
268                    )
269
270                    if is_type_change:
271                        column_to_alter = action.this
272                        new_type = action.args["dtype"]
273                        temp_column_name_str = f"{column_to_alter.name}__{random_id(short=True)}"
274                        temp_column_name = exp.to_identifier(temp_column_name_str)
275
276                        logger.info(
277                            "Applying workaround for column '%s' on table '%s' to change type to '%s'.",
278                            column_to_alter.sql(),
279                            table_name.sql(),
280                            new_type.sql(),
281                        )
282
283                        # Step 1: Add a temporary column.
284                        add_column_expr = exp.Alter(
285                            this=table_name_without_catalog.copy(),
286                            kind="TABLE",
287                            actions=[
288                                exp.ColumnDef(this=temp_column_name.copy(), kind=new_type.copy())
289                            ],
290                        )
291                        add_sql = self._to_sql(add_column_expr)
292                        self.execute(add_sql)
293
294                        # Step 2: Copy and cast data.
295                        update_sql = self._to_sql(
296                            exp.Update(
297                                this=table_name_without_catalog.copy(),
298                                expressions=[
299                                    exp.EQ(
300                                        this=temp_column_name.copy(),
301                                        expression=exp.Cast(
302                                            this=column_to_alter.copy(), to=new_type.copy()
303                                        ),
304                                    )
305                                ],
306                            )
307                        )
308                        self.execute(update_sql)
309
310                        # Step 3: Drop the original column.
311                        drop_sql = self._to_sql(
312                            exp.Alter(
313                                this=table_name_without_catalog.copy(),
314                                kind="TABLE",
315                                actions=[exp.Drop(this=column_to_alter.copy(), kind="COLUMN")],
316                            )
317                        )
318                        self.execute(drop_sql)
319
320                        # Step 4: Rename the temporary column.
321                        old_name_qualified = f"{table_name_without_catalog.sql(dialect=self.dialect)}.{temp_column_name.sql(dialect=self.dialect)}"
322                        new_name_unquoted = column_to_alter.sql(
323                            dialect=self.dialect, identify=False
324                        )
325                        rename_sql = f"EXEC sp_rename '{old_name_qualified}', '{new_name_unquoted}', 'COLUMN'"
326                        self.execute(rename_sql)
327                    else:
328                        # For other alterations, execute directly.
329                        direct_alter_expr = exp.Alter(
330                            this=table_name_without_catalog.copy(), kind="TABLE", actions=[action]
331                        )
332                        self.execute(direct_alter_expr)
333
334
335class FabricHttpClient:
336    def __init__(self, tenant_id: str, workspace_id: str, client_id: str, client_secret: str):
337        self.tenant_id = tenant_id
338        self.client_id = client_id
339        self.client_secret = client_secret
340        self.workspace_id = workspace_id
341
342    def create_warehouse(
343        self, warehouse_name: str, if_not_exists: bool = True, attempt: int = 0
344    ) -> None:
345        """Create a catalog (warehouse) in Microsoft Fabric via REST API."""
346
347        # attempt count is arbitrary, it essentially equates to 5 minutes of 30 second waits
348        if attempt > 10:
349            raise SQLMeshError(
350                f"Gave up waiting for Fabric warehouse {warehouse_name} to become available"
351            )
352
353        logger.info(f"Creating Fabric warehouse: {warehouse_name}")
354
355        request_data = {
356            "displayName": warehouse_name,
357            "description": f"Warehouse created by SQLMesh: {warehouse_name}",
358        }
359
360        response = self.session.post(self._endpoint_url("warehouses"), json=request_data)
361
362        if (
363            if_not_exists
364            and response.status_code == 400
365            and (errorCode := response.json().get("errorCode", None))
366        ):
367            if errorCode == "ItemDisplayNameAlreadyInUse":
368                logger.warning(f"Fabric warehouse {warehouse_name} already exists")
369                return
370            if errorCode == "ItemDisplayNameNotAvailableYet":
371                logger.warning(f"Fabric warehouse {warehouse_name} is still spinning up; waiting")
372                # Fabric error message is something like:
373                #  - "Requested 'circleci_51d7087e__dev' is not available yet and is expected to become available in the upcoming minutes."
374                # This seems to happen if a catalog is dropped and then a new one with the same name is immediately created.
375                # There appears to be some delayed async process on the Fabric side that actually drops the warehouses and frees up the names to be used again
376                time.sleep(30)
377                return self.create_warehouse(
378                    warehouse_name=warehouse_name, if_not_exists=if_not_exists, attempt=attempt + 1
379                )
380
381        try:
382            response.raise_for_status()
383        except:
384            # the important information to actually debug anything is in the response body which Requests never prints
385            logger.exception(
386                f"Failed to create warehouse {warehouse_name}. status: {response.status_code}, body: {response.text}"
387            )
388            raise
389
390        # Handle direct success (201) or async creation (202)
391        if response.status_code == 201:
392            logger.info(f"Successfully created Fabric warehouse: {warehouse_name}")
393            return
394
395        if response.status_code == 202 and (location_header := response.headers.get("location")):
396            logger.info(f"Warehouse creation initiated for: {warehouse_name}")
397            self._wait_for_completion(location_header, warehouse_name)
398            logger.info(f"Successfully created Fabric warehouse: {warehouse_name}")
399        else:
400            logger.error(f"Unexpected response from Fabric API: {response}\n{response.text}")
401            raise SQLMeshError(f"Unable to create warehouse: {response}")
402
403    def delete_warehouse(self, warehouse_name: str, if_exists: bool = True) -> None:
404        """Drop a catalog (warehouse) in Microsoft Fabric via REST API."""
405        logger.info(f"Deleting Fabric warehouse: {warehouse_name}")
406
407        # Get the warehouse ID by listing warehouses
408        # TODO: handle continuationUri for pagination, ref: https://learn.microsoft.com/en-us/rest/api/fabric/warehouse/items/list-warehouses?tabs=HTTP#warehouses
409        response = self.session.get(self._endpoint_url("warehouses"))
410        response.raise_for_status()
411
412        warehouse_name_to_id = {
413            warehouse.get("displayName"): warehouse.get("id")
414            for warehouse in response.json().get("value", [])
415        }
416
417        warehouse_id = warehouse_name_to_id.get(warehouse_name, None)
418
419        if not warehouse_id:
420            logger.warning(
421                f"Fabric warehouse does not exist: {warehouse_name}\n(available warehouses: {', '.join(warehouse_name_to_id)})"
422            )
423            if if_exists:
424                return
425
426            raise SQLMeshError(
427                f"Unable to delete Fabric warehouse {warehouse_name} as it doesnt exist"
428            )
429
430        # Delete the warehouse by ID
431        response = self.session.delete(self._endpoint_url(f"warehouses/{warehouse_id}"))
432        response.raise_for_status()
433
434        logger.info(f"Successfully deleted Fabric warehouse: {warehouse_name}")
435
436    @cached_property
437    def session(self) -> requests.Session:
438        s = requests.Session()
439
440        access_token = self._get_access_token()
441        s.headers.update({"Authorization": f"Bearer {access_token}"})
442
443        return s
444
445    def _endpoint_url(self, endpoint: str) -> str:
446        if endpoint.startswith("/"):
447            endpoint = endpoint[1:]
448
449        return f"https://api.fabric.microsoft.com/v1/workspaces/{self.workspace_id}/{endpoint}"
450
451    def _get_access_token(self) -> str:
452        """Get access token using Service Principal authentication."""
453
454        # Use Azure AD OAuth2 token endpoint
455        token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
456
457        data = {
458            "grant_type": "client_credentials",
459            "client_id": self.client_id,
460            "client_secret": self.client_secret,
461            "scope": "https://api.fabric.microsoft.com/.default",
462        }
463
464        response = requests.post(token_url, data=data)
465        response.raise_for_status()
466        token_data = response.json()
467        return token_data["access_token"]
468
469    def _wait_for_completion(self, location_url: str, operation_name: str) -> None:
470        """Poll the operation status until completion."""
471
472        @retry(
473            wait=wait_exponential(multiplier=1, min=1, max=30),
474            stop=stop_after_attempt(20),
475            retry=retry_if_result(lambda result: result not in ["Succeeded", "Failed"]),
476        )
477        def _poll() -> str:
478            response = self.session.get(location_url)
479            response.raise_for_status()
480
481            result = response.json()
482            status = result.get("status", "Unknown")
483
484            logger.debug(f"Operation {operation_name} status: {status}")
485
486            if status == "Failed":
487                error_msg = result.get("error", {}).get("message", "Unknown error")
488                raise SQLMeshError(f"Operation {operation_name} failed: {error_msg}")
489            elif status in ["InProgress", "Running"]:
490                logger.debug(f"Operation {operation_name} still in progress...")
491            elif status not in ["Succeeded"]:
492                logger.warning(f"Unknown status '{status}' for operation {operation_name}")
493
494            return status
495
496        final_status = _poll()
497        if final_status != "Succeeded":
498            raise SQLMeshError(f"Operation {operation_name} completed with status: {final_status}")