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 api_client(self) -> FabricHttpClient:
 63        # the requests Session is not guaranteed to be threadsafe
 64        # so we create a http client per thread on demand
 65        if existing_client := self._connection_pool.get_attribute("api_client"):
 66            return existing_client
 67
 68        tenant_id: t.Optional[str] = self._extra_config.get("tenant_id")
 69        workspace_id: t.Optional[str] = self._extra_config.get("workspace_id")
 70        client_id: t.Optional[str] = self._extra_config.get("user")
 71        client_secret: t.Optional[str] = self._extra_config.get("password")
 72
 73        if not tenant_id or not client_id or not client_secret:
 74            raise SQLMeshError(
 75                "Service Principal authentication requires tenant_id, client_id, and client_secret "
 76                "in the Fabric connection configuration"
 77            )
 78
 79        if not workspace_id:
 80            raise SQLMeshError(
 81                "Fabric requires the workspace_id to be configured in the connection configuration to create / drop catalogs"
 82            )
 83
 84        client = FabricHttpClient(
 85            tenant_id=tenant_id,
 86            workspace_id=workspace_id,
 87            client_id=client_id,
 88            client_secret=client_secret,
 89        )
 90
 91        self._connection_pool.set_attribute("api_client", client)
 92        return client
 93
 94    def _create_catalog(self, catalog_name: exp.Identifier) -> None:
 95        """Create a catalog (warehouse) in Microsoft Fabric via REST API."""
 96        warehouse_name = catalog_name.sql(dialect=self.dialect, identify=False)
 97        logger.info(f"Creating Fabric warehouse: {warehouse_name}")
 98
 99        self.api_client.create_warehouse(warehouse_name)
100
101    def _drop_catalog(self, catalog_name: exp.Identifier) -> None:
102        """Drop a catalog (warehouse) in Microsoft Fabric via REST API."""
103        warehouse_name = catalog_name.sql(dialect=self.dialect, identify=False)
104        current_catalog = self.get_current_catalog()
105
106        logger.info(f"Deleting Fabric warehouse: {warehouse_name}")
107        self.api_client.delete_warehouse(warehouse_name)
108
109        if warehouse_name == current_catalog:
110            # Somewhere around 2025-09-08, Fabric started validating the "Database=" connection argument and throwing 'Authentication failed' if the database doesnt exist
111            # In addition, set_current_catalog() is implemented using a threadlocal variable "target_catalog"
112            # So, when we drop a warehouse, and there are still threads with "target_catalog" set to reference it, any operations on those threads
113            # that use an either use an existing connection pointing to this warehouse or trigger a new connection
114            # will fail with an 'Authentication Failed' error unless we close all connections here, which also clears all the threadlocal data
115            self.close()
116
117    def set_current_catalog(self, catalog_name: str) -> None:
118        """
119        Set the current catalog for Microsoft Fabric connections.
120
121        Override to handle Fabric's stateless session limitation where USE statements
122        don't persist across queries. Instead, we close existing connections and
123        recreate them with the new catalog in the connection configuration.
124
125        Args:
126            catalog_name: The name of the catalog (warehouse) to switch to
127
128        Note:
129            Fabric doesn't support catalog switching via USE statements because each
130            statement runs as an independent session. This method works around this
131            limitation by updating the connection pool with new catalog configuration.
132
133        See:
134            https://learn.microsoft.com/en-us/fabric/data-warehouse/sql-query-editor#limitations
135        """
136        current_catalog = self.get_current_catalog()
137
138        # If already using the requested catalog, do nothing
139        if current_catalog and current_catalog == catalog_name:
140            logger.debug(f"Already using catalog '{catalog_name}', no action needed")
141            return
142
143        logger.info(f"Switching from catalog '{current_catalog}' to '{catalog_name}'")
144
145        # commit the transaction before closing the connection to help prevent errors like:
146        # > Snapshot isolation transaction failed in database because the object accessed by the statement has been modified by a
147        # > DDL statement in another concurrent transaction since the start of this transaction
148        # on subsequent queries in the new connection
149        self._connection_pool.commit()
150
151        # note: we call close() on the connection pool instead of self.close() because self.close() calls close_all()
152        # on the connection pool but we just want to close the connection for this thread
153        self._connection_pool.close()
154        self._target_catalog = catalog_name  # new connections will use this catalog
155
156        catalog_after_switch = self.get_current_catalog()
157
158        if catalog_after_switch != catalog_name:
159            # We need to raise an error if the catalog switch failed to prevent the operation that needed the catalog switch from being run against the wrong catalog
160            raise SQLMeshError(
161                f"Unable to switch catalog to {catalog_name}, catalog ended up as {catalog_after_switch}"
162            )
163
164    def alter_table(
165        self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]]
166    ) -> None:
167        """
168        Applies alter expressions to a table. Fabric has limited support for ALTER TABLE,
169        so this method implements a workaround for column type changes.
170        This method is self-contained and sets its own catalog context.
171        """
172        if not alter_expressions:
173            return
174
175        # Get the target table from the first expression to determine the correct catalog.
176        first_op = alter_expressions[0]
177        expression = first_op.expression if isinstance(first_op, TableAlterOperation) else first_op
178        if not isinstance(expression, exp.Alter) or not expression.this.catalog:
179            # Fallback for unexpected scenarios
180            logger.warning(
181                "Could not determine catalog from alter expression, executing with current context."
182            )
183            super().alter_table(alter_expressions)
184            return
185
186        target_catalog = expression.this.catalog
187        self.set_current_catalog(target_catalog)
188
189        with self.transaction():
190            for op in alter_expressions:
191                expression = op.expression if isinstance(op, TableAlterOperation) else op
192
193                if not isinstance(expression, exp.Alter):
194                    self.execute(expression)
195                    continue
196
197                for action in expression.actions:
198                    table_name = expression.this
199
200                    table_name_without_catalog = table_name.copy()
201                    table_name_without_catalog.set("catalog", None)
202
203                    is_type_change = isinstance(action, exp.AlterColumn) and action.args.get(
204                        "dtype"
205                    )
206
207                    if is_type_change:
208                        column_to_alter = action.this
209                        new_type = action.args["dtype"]
210                        temp_column_name_str = f"{column_to_alter.name}__{random_id(short=True)}"
211                        temp_column_name = exp.to_identifier(temp_column_name_str)
212
213                        logger.info(
214                            "Applying workaround for column '%s' on table '%s' to change type to '%s'.",
215                            column_to_alter.sql(),
216                            table_name.sql(),
217                            new_type.sql(),
218                        )
219
220                        # Step 1: Add a temporary column.
221                        add_column_expr = exp.Alter(
222                            this=table_name_without_catalog.copy(),
223                            kind="TABLE",
224                            actions=[
225                                exp.ColumnDef(this=temp_column_name.copy(), kind=new_type.copy())
226                            ],
227                        )
228                        add_sql = self._to_sql(add_column_expr)
229                        self.execute(add_sql)
230
231                        # Step 2: Copy and cast data.
232                        update_sql = self._to_sql(
233                            exp.Update(
234                                this=table_name_without_catalog.copy(),
235                                expressions=[
236                                    exp.EQ(
237                                        this=temp_column_name.copy(),
238                                        expression=exp.Cast(
239                                            this=column_to_alter.copy(), to=new_type.copy()
240                                        ),
241                                    )
242                                ],
243                            )
244                        )
245                        self.execute(update_sql)
246
247                        # Step 3: Drop the original column.
248                        drop_sql = self._to_sql(
249                            exp.Alter(
250                                this=table_name_without_catalog.copy(),
251                                kind="TABLE",
252                                actions=[exp.Drop(this=column_to_alter.copy(), kind="COLUMN")],
253                            )
254                        )
255                        self.execute(drop_sql)
256
257                        # Step 4: Rename the temporary column.
258                        old_name_qualified = f"{table_name_without_catalog.sql(dialect=self.dialect)}.{temp_column_name.sql(dialect=self.dialect)}"
259                        new_name_unquoted = column_to_alter.sql(
260                            dialect=self.dialect, identify=False
261                        )
262                        rename_sql = f"EXEC sp_rename '{old_name_qualified}', '{new_name_unquoted}', 'COLUMN'"
263                        self.execute(rename_sql)
264                    else:
265                        # For other alterations, execute directly.
266                        direct_alter_expr = exp.Alter(
267                            this=table_name_without_catalog.copy(), kind="TABLE", actions=[action]
268                        )
269                        self.execute(direct_alter_expr)
270
271
272class FabricHttpClient:
273    def __init__(self, tenant_id: str, workspace_id: str, client_id: str, client_secret: str):
274        self.tenant_id = tenant_id
275        self.client_id = client_id
276        self.client_secret = client_secret
277        self.workspace_id = workspace_id
278
279    def create_warehouse(
280        self, warehouse_name: str, if_not_exists: bool = True, attempt: int = 0
281    ) -> None:
282        """Create a catalog (warehouse) in Microsoft Fabric via REST API."""
283
284        # attempt count is arbitrary, it essentially equates to 5 minutes of 30 second waits
285        if attempt > 10:
286            raise SQLMeshError(
287                f"Gave up waiting for Fabric warehouse {warehouse_name} to become available"
288            )
289
290        logger.info(f"Creating Fabric warehouse: {warehouse_name}")
291
292        request_data = {
293            "displayName": warehouse_name,
294            "description": f"Warehouse created by SQLMesh: {warehouse_name}",
295        }
296
297        response = self.session.post(self._endpoint_url("warehouses"), json=request_data)
298
299        if (
300            if_not_exists
301            and response.status_code == 400
302            and (errorCode := response.json().get("errorCode", None))
303        ):
304            if errorCode == "ItemDisplayNameAlreadyInUse":
305                logger.warning(f"Fabric warehouse {warehouse_name} already exists")
306                return
307            if errorCode == "ItemDisplayNameNotAvailableYet":
308                logger.warning(f"Fabric warehouse {warehouse_name} is still spinning up; waiting")
309                # Fabric error message is something like:
310                #  - "Requested 'circleci_51d7087e__dev' is not available yet and is expected to become available in the upcoming minutes."
311                # This seems to happen if a catalog is dropped and then a new one with the same name is immediately created.
312                # 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
313                time.sleep(30)
314                return self.create_warehouse(
315                    warehouse_name=warehouse_name, if_not_exists=if_not_exists, attempt=attempt + 1
316                )
317
318        try:
319            response.raise_for_status()
320        except:
321            # the important information to actually debug anything is in the response body which Requests never prints
322            logger.exception(
323                f"Failed to create warehouse {warehouse_name}. status: {response.status_code}, body: {response.text}"
324            )
325            raise
326
327        # Handle direct success (201) or async creation (202)
328        if response.status_code == 201:
329            logger.info(f"Successfully created Fabric warehouse: {warehouse_name}")
330            return
331
332        if response.status_code == 202 and (location_header := response.headers.get("location")):
333            logger.info(f"Warehouse creation initiated for: {warehouse_name}")
334            self._wait_for_completion(location_header, warehouse_name)
335            logger.info(f"Successfully created Fabric warehouse: {warehouse_name}")
336        else:
337            logger.error(f"Unexpected response from Fabric API: {response}\n{response.text}")
338            raise SQLMeshError(f"Unable to create warehouse: {response}")
339
340    def delete_warehouse(self, warehouse_name: str, if_exists: bool = True) -> None:
341        """Drop a catalog (warehouse) in Microsoft Fabric via REST API."""
342        logger.info(f"Deleting Fabric warehouse: {warehouse_name}")
343
344        # Get the warehouse ID by listing warehouses
345        # TODO: handle continuationUri for pagination, ref: https://learn.microsoft.com/en-us/rest/api/fabric/warehouse/items/list-warehouses?tabs=HTTP#warehouses
346        response = self.session.get(self._endpoint_url("warehouses"))
347        response.raise_for_status()
348
349        warehouse_name_to_id = {
350            warehouse.get("displayName"): warehouse.get("id")
351            for warehouse in response.json().get("value", [])
352        }
353
354        warehouse_id = warehouse_name_to_id.get(warehouse_name, None)
355
356        if not warehouse_id:
357            logger.warning(
358                f"Fabric warehouse does not exist: {warehouse_name}\n(available warehouses: {', '.join(warehouse_name_to_id)})"
359            )
360            if if_exists:
361                return
362
363            raise SQLMeshError(
364                f"Unable to delete Fabric warehouse {warehouse_name} as it doesnt exist"
365            )
366
367        # Delete the warehouse by ID
368        response = self.session.delete(self._endpoint_url(f"warehouses/{warehouse_id}"))
369        response.raise_for_status()
370
371        logger.info(f"Successfully deleted Fabric warehouse: {warehouse_name}")
372
373    @cached_property
374    def session(self) -> requests.Session:
375        s = requests.Session()
376
377        access_token = self._get_access_token()
378        s.headers.update({"Authorization": f"Bearer {access_token}"})
379
380        return s
381
382    def _endpoint_url(self, endpoint: str) -> str:
383        if endpoint.startswith("/"):
384            endpoint = endpoint[1:]
385
386        return f"https://api.fabric.microsoft.com/v1/workspaces/{self.workspace_id}/{endpoint}"
387
388    def _get_access_token(self) -> str:
389        """Get access token using Service Principal authentication."""
390
391        # Use Azure AD OAuth2 token endpoint
392        token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
393
394        data = {
395            "grant_type": "client_credentials",
396            "client_id": self.client_id,
397            "client_secret": self.client_secret,
398            "scope": "https://api.fabric.microsoft.com/.default",
399        }
400
401        response = requests.post(token_url, data=data)
402        response.raise_for_status()
403        token_data = response.json()
404        return token_data["access_token"]
405
406    def _wait_for_completion(self, location_url: str, operation_name: str) -> None:
407        """Poll the operation status until completion."""
408
409        @retry(
410            wait=wait_exponential(multiplier=1, min=1, max=30),
411            stop=stop_after_attempt(20),
412            retry=retry_if_result(lambda result: result not in ["Succeeded", "Failed"]),
413        )
414        def _poll() -> str:
415            response = self.session.get(location_url)
416            response.raise_for_status()
417
418            result = response.json()
419            status = result.get("status", "Unknown")
420
421            logger.debug(f"Operation {operation_name} status: {status}")
422
423            if status == "Failed":
424                error_msg = result.get("error", {}).get("message", "Unknown error")
425                raise SQLMeshError(f"Operation {operation_name} failed: {error_msg}")
426            elif status in ["InProgress", "Running"]:
427                logger.debug(f"Operation {operation_name} still in progress...")
428            elif status not in ["Succeeded"]:
429                logger.warning(f"Unknown status '{status}' for operation {operation_name}")
430
431            return status
432
433        final_status = _poll()
434        if final_status != "Succeeded":
435            raise SQLMeshError(f"Operation {operation_name} completed with status: {final_status}")