feat(duckdb): add file systems configuration option (#4778)

Author: mrjsj

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.
+> 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. 

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 @@ def init(cursor: duckdb.DuckDBPyConnection) -> None:
                             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()
             ):

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 @@ def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
             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,

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",
     }
 

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
 

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>",
     ]