Repository: TobikoData/sqlmesh
Pull request #4794 (2025-06-23 21:31:55)

Fix: Create a transaction and a session when migrating a snapshot
The statements that run as part of the migration needs to be a part of the same transaction + session.


Patch:
diff --git a/docs/integrations/engines/duckdb.md b/docs/integrations/engines/duckdb.md
index 7aef7def..19d8d6e1 100644
--- a/docs/integrations/engines/duckdb.md
+++ b/docs/integrations/engines/duckdb.md
@@ -18,6 +18,7 @@
 | `extensions`       | Extension to load into duckdb. Only autoloadable extensions are supported.                                                                                                                                                                      |  list  |    N     |
 | `connector_config` | Configuration to pass into the duckdb connector.                                                                                                                                                                                                |  dict  |    N     |
 | `secrets`   | Configuration for authenticating external sources (e.g., S3) using DuckDB secrets.                                                                                                                                                                                                |  dict  |    N     |
+| `filesystems`   | Configuration for registering `fsspec` filesystems to the DuckDB connection.                                                                                                                                                                                                |  dict  |    N     |
 
 #### DuckDB Catalogs Example
 
@@ -256,4 +257,36 @@ After configuring the secrets, you can directly reference S3 paths in your catal
 
 Refer to the official DuckDB documentation for the full list of [supported S3 secret parameters](https://duckdb.org/docs/stable/extensions/httpfs/s3api.html#overview-of-s3-secret-parameters) and for more information on the [Secrets Manager configuration](https://duckdb.org/docs/configuration/secrets_manager.html).
 
-> Note: Loading credentials at runtime using `load_aws_credentials()` or similar deprecated functions may fail when using SQLMesh.
\ No newline at end of file
+> Note: Loading credentials at runtime using `load_aws_credentials()` or similar deprecated functions may fail when using SQLMesh.
+
+##### File system configuration example for Microsoft Onelake
+
+The `filesystems` accepts a list of file systems to register in the DuckDB connection. This is especially useful for Azure Storage Accounts, as it adds write support for DuckDB which is not natively supported by DuckDB (yet).
+
+
+=== "YAML"
+
+    ```yaml linenums="1"
+    gateways:
+      ducklake:
+        connection:
+          type: duckdb
+          catalogs:
+            ducklake:
+            type: ducklake
+            path: myducklakecatalog.duckdb
+            data_path: abfs://MyFabricWorkspace/MyFabricLakehouse.Lakehouse/Files/DuckLake.Files
+        extensions:
+          - ducklake
+        filesystems:
+          - fs: abfs
+            account_name: onelake
+            account_host: onelake.blob.fabric.microsoft.com
+            client_id: {{ env_var('AZURE_CLIENT_ID') }}
+            client_secret: {{ env_var('AZURE_CLIENT_SECRET') }}
+            tenant_id: {{ env_var('AZURE_TENANT_ID') }}
+            # anon: False # To use azure.identity.DefaultAzureCredential authentication 
+    ```
+
+
+Refer to the documentation for `fsspec` [fsspec.filesystem](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.filesystem) and `adlfs` [adlfs.AzureBlobFileSystem](https://fsspec.github.io/adlfs/api/#api-reference) for a full list of storage options. 
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index a669471d..2c897cd8 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -266,6 +266,7 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
         extensions: A list of autoloadable extensions to load.
         connector_config: A dictionary of configuration to pass into the duckdb connector.
         secrets: A list of dictionaries used to generate DuckDB secrets for authenticating with external services (e.g. S3).
+        filesystems: A list of dictionaries used to register `fsspec` filesystems to the DuckDB cursor.
         concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
         register_comments: Whether or not to register model comments with the SQL engine.
         pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
@@ -277,6 +278,7 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
     extensions: t.List[t.Union[str, t.Dict[str, t.Any]]] = []
     connector_config: t.Dict[str, t.Any] = {}
     secrets: t.List[t.Dict[str, t.Any]] = []
+    filesystems: t.List[t.Dict[str, t.Any]] = []
 
     concurrent_tasks: int = 1
     register_comments: bool = True
@@ -371,6 +373,15 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
                             except Exception as e:
                                 raise ConfigError(f"Failed to create secret: {e}")
 
+            if self.filesystems:
+                from fsspec import filesystem  # type: ignore
+
+                for file_system in self.filesystems:
+                    options = file_system.copy()
+                    fs = options.pop("fs")
+                    fs = filesystem(fs, **options)
+                    cursor.register_filesystem(fs)
+
             for i, (alias, path_options) in enumerate(
                 (getattr(self, "catalogs", None) or {}).items()
             ):
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index f2042583..d3374863 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -911,39 +911,40 @@ class SnapshotEvaluator:
         ):
             return
 
+        deployability_index = DeployabilityIndex.all_deployable()
+        render_kwargs: t.Dict[str, t.Any] = dict(
+            engine_adapter=adapter,
+            snapshots=parent_snapshots_by_name(snapshot, snapshots),
+            runtime_stage=RuntimeStage.CREATING,
+            deployability_index=deployability_index,
+        )
         target_table_name = snapshot.table_name()
-        if adapter.table_exists(target_table_name):
-            evaluation_strategy = _evaluation_strategy(snapshot, adapter)
-            tmp_table_name = snapshot.table_name(is_deployable=False)
-            logger.info(
-                "Migrating table schema from '%s' to '%s'",
-                tmp_table_name,
-                target_table_name,
-            )
-            evaluation_strategy.migrate(
-                target_table_name=target_table_name,
-                source_table_name=tmp_table_name,
-                snapshot=snapshot,
-                snapshots=parent_snapshots_by_name(snapshot, snapshots),
-                allow_destructive_snapshots=allow_destructive_snapshots,
-            )
-        else:
-            logger.info(
-                "Creating table '%s' for the snapshot of the forward-only model %s",
-                target_table_name,
-                snapshot.snapshot_id,
-            )
-            deployability_index = DeployabilityIndex.all_deployable()
-            render_kwargs: t.Dict[str, t.Any] = dict(
-                engine_adapter=adapter,
-                snapshots=parent_snapshots_by_name(snapshot, snapshots),
-                runtime_stage=RuntimeStage.CREATING,
-                deployability_index=deployability_index,
-            )
-            with (
-                adapter.transaction(),
-                adapter.session(snapshot.model.render_session_properties(**render_kwargs)),
-            ):
+
+        with (
+            adapter.transaction(),
+            adapter.session(snapshot.model.render_session_properties(**render_kwargs)),
+        ):
+            if adapter.table_exists(target_table_name):
+                evaluation_strategy = _evaluation_strategy(snapshot, adapter)
+                tmp_table_name = snapshot.table_name(is_deployable=False)
+                logger.info(
+                    "Migrating table schema from '%s' to '%s'",
+                    tmp_table_name,
+                    target_table_name,
+                )
+                evaluation_strategy.migrate(
+                    target_table_name=target_table_name,
+                    source_table_name=tmp_table_name,
+                    snapshot=snapshot,
+                    snapshots=parent_snapshots_by_name(snapshot, snapshots),
+                    allow_destructive_snapshots=allow_destructive_snapshots,
+                )
+            else:
+                logger.info(
+                    "Creating table '%s' for the snapshot of the forward-only model %s",
+                    target_table_name,
+                    snapshot.snapshot_id,
+                )
                 self._execute_create(
                     snapshot=snapshot,
                     table_name=target_table_name,
diff --git a/sqlmesh/dbt/target.py b/sqlmesh/dbt/target.py
index 05985d87..94f4c988 100644
--- a/sqlmesh/dbt/target.py
+++ b/sqlmesh/dbt/target.py
@@ -138,6 +138,7 @@ class DuckDbConfig(TargetConfig):
         extensions: A list of autoloadable extensions to load.
         settings: A dictionary of settings to pass into the duckdb connector.
         secrets: A list of secrets to pass to the secret manager in the duckdb connector.
+        filesystems: A list of `fsspec` filesystems to register in the duckdb connection.
     """
 
     type: t.Literal["duckdb"] = "duckdb"
@@ -147,6 +148,7 @@ class DuckDbConfig(TargetConfig):
     extensions: t.Optional[t.List[str]] = None
     settings: t.Optional[t.Dict[str, t.Any]] = None
     secrets: t.Optional[t.List[t.Dict[str, t.Any]]] = None
+    filesystems: t.Optional[t.List[t.Dict[str, t.Any]]] = None
 
     @model_validator(mode="before")
     def validate_authentication(cls, data: t.Any) -> t.Any:
@@ -182,6 +184,8 @@ class DuckDbConfig(TargetConfig):
             kwargs["connector_config"] = self.settings
         if self.secrets is not None:
             kwargs["secrets"] = self.secrets
+        if self.filesystems is not None:
+            kwargs["filesystems"] = self.filesystems
         return DuckDBConnectionConfig(
             database=self.path,
             concurrent_tasks=1,


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 47e063c5..b3457345 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -562,6 +562,7 @@ def test_connection_config_serialization():
         "pretty_sql": False,
         "connector_config": {},
         "secrets": [],
+        "filesystems": [],
         "database": "my_db",
     }
     assert serialized["default_test_connection"] == {
@@ -573,6 +574,7 @@ def test_connection_config_serialization():
         "pretty_sql": False,
         "connector_config": {},
         "secrets": [],
+        "filesystems": [],
         "database": "my_test_db",
     }
 
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 90aba1c3..0d7df3d7 100644
--- a/tests/core/test_connection_config.py
+++ b/tests/core/test_connection_config.py
@@ -437,9 +437,20 @@ def test_duckdb(make_config):
                 "secret": "aws_secret",
             }
         ],
+        filesystems=[
+            {
+                "protocol": "abfs",
+                "storage_options": {
+                    "account_name": "onelake",
+                    "account_host": "onelake.blob.fabric.microsoft.com",
+                    "anon": False,
+                },
+            }
+        ],
     )
     assert config.connector_config
     assert config.secrets
+    assert config.filesystems
     assert isinstance(config, DuckDBConnectionConfig)
     assert not config.is_recommended_for_state_sync
 
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index d572fc7e..d131e6aa 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -54,7 +54,7 @@ from sqlmesh.core.snapshot import (
     SnapshotTableCleanupTask,
 )
 from sqlmesh.core.snapshot.definition import to_view_mapping
-from sqlmesh.core.snapshot.evaluator import CustomMaterialization
+from sqlmesh.core.snapshot.evaluator import CustomMaterialization, SnapshotCreationFailedError
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.date import to_timestamp
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, DestructiveChangeError
@@ -92,13 +92,16 @@ def date_kwargs() -> t.Dict[str, str]:
 
 @pytest.fixture
 def adapter_mock(mocker: MockerFixture):
+    def mock_exit(self, exc_type, exc_value, traceback):
+        pass
+
     transaction_mock = mocker.Mock()
     transaction_mock.__enter__ = mocker.Mock()
-    transaction_mock.__exit__ = mocker.Mock()
+    transaction_mock.__exit__ = mock_exit
 
     session_mock = mocker.Mock()
     session_mock.__enter__ = mocker.Mock()
-    session_mock.__exit__ = mocker.Mock()
+    session_mock.__exit__ = mock_exit
 
     adapter_mock = mocker.Mock()
     adapter_mock.transaction.return_value = transaction_mock
@@ -1160,6 +1163,7 @@ def test_migrate(mocker: MockerFixture, make_snapshot):
     cursor_mock = mocker.Mock()
     connection_mock.cursor.return_value = cursor_mock
     adapter = EngineAdapter(lambda: connection_mock, "")
+    session_spy = mocker.spy(adapter, "session")
 
     current_table = "sqlmesh__test_schema.test_schema__test_model__1"
 
@@ -1201,6 +1205,8 @@ def test_migrate(mocker: MockerFixture, make_snapshot):
         ]
     )
 
+    session_spy.assert_called_once()
+
 
 def test_migrate_missing_table(mocker: MockerFixture, make_snapshot):
     connection_mock = mocker.NonCallableMock()
@@ -1596,7 +1602,8 @@ def test_drop_clone_in_dev_when_migration_fails(mocker: MockerFixture, adapter_m
         ),
     ]
 
-    evaluator.create([snapshot], {})
+    with pytest.raises(SnapshotCreationFailedError):
+        evaluator.create([snapshot], {})
 
     adapter_mock.clone_table.assert_called_once_with(
         f"sqlmesh__test_schema.test_schema__test_model__{snapshot.version}__dev",
@@ -2537,7 +2544,9 @@ def test_create_seed_on_error(mocker: MockerFixture, adapter_mock, make_snapshot
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
     evaluator = SnapshotEvaluator(adapter_mock)
-    evaluator.create([snapshot], {})
+
+    with pytest.raises(SnapshotCreationFailedError):
+        evaluator.create([snapshot], {})
 
     adapter_mock.replace_query.assert_called_once_with(
         f"sqlmesh__db.db__seed__{snapshot.version}",
diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py
index f606e574..6507f071 100644
--- a/tests/integrations/jupyter/test_magics.py
+++ b/tests/integrations/jupyter/test_magics.py
@@ -745,16 +745,16 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
         "Models: 20",
         "Macros: 8",
         "",
-        "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None",
-        "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None",
+        "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []",
+        "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []",
         "Data warehouse connection succeeded",
     ]
     assert get_all_html_output(output) == [
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Models: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">20</span></pre>",
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Macros: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span></pre>",
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>",
-        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span></pre>',
-        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span></pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span></pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span></pre>',
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Data warehouse connection <span style=\"color: #008000; text-decoration-color: #008000\">succeeded</span></pre>",
     ]