Merge pull request #1221 from matteomorlack/feature/register-filesystems

feat: register filesystems

Author: Mause

README.md

@@ -5,19 +5,20 @@
 Basic SQLAlchemy driver for [DuckDB](https://duckdb.org/)
 
 <!--ts-->
-* [duckdb_engine](#duckdb_engine)
-   * [Installation](#installation)
-   * [Usage](#usage)
-   * [Usage in IPython/Jupyter](#usage-in-ipythonjupyter)
-   * [Configuration](#configuration)
-   * [How to register a pandas DataFrame](#how-to-register-a-pandas-dataframe)
-   * [Things to keep in mind](#things-to-keep-in-mind)
-      * [Auto-incrementing ID columns](#auto-incrementing-id-columns)
-      * [Pandas read_sql() chunksize](#pandas-read_sql-chunksize)
-      * [Unsigned integer support](#unsigned-integer-support)
-   * [Alembic Integration](#alembic-integration)
-   * [Preloading extensions (experimental)](#preloading-extensions-experimental)
-   * [The name](#the-name)
+- [duckdb\_engine](#duckdb_engine)
+  - [Installation](#installation)
+  - [Usage](#usage)
+  - [Usage in IPython/Jupyter](#usage-in-ipythonjupyter)
+  - [Configuration](#configuration)
+  - [How to register a pandas DataFrame](#how-to-register-a-pandas-dataframe)
+  - [Things to keep in mind](#things-to-keep-in-mind)
+    - [Auto-incrementing ID columns](#auto-incrementing-id-columns)
+    - [Pandas `read_sql()` chunksize](#pandas-read_sql-chunksize)
+    - [Unsigned integer support](#unsigned-integer-support)
+  - [Alembic Integration](#alembic-integration)
+  - [Preloading extensions (experimental)](#preloading-extensions-experimental)
+  - [Registering Filesystems](#registering-filesystems)
+  - [The name](#the-name)
 
 <!-- Created by https://github.com/ekalinin/github-markdown-toc -->
 <!-- Added by: me, at: Wed 20 Sep 2023 12:44:27 AWST -->
@@ -180,6 +181,24 @@ create_engine(
 )
 ```
 
+## Registering Filesystems
+
+> DuckDB allows registering filesystems from [fsspec](https://filesystem-spec.readthedocs.io/), see [documentation](https://duckdb.org/docs/guides/python/filesystems.html) for more information.
+
+Support is provided under `connect_args` parameter
+
+```python
+from sqlalchemy import create_engine
+from fsspec import filesystem
+
+create_engine(
+    'duckdb:///:memory:',
+    connect_args={
+        'register_filesystems': [filesystem('gcs')],
+    }
+)
+```
+
 ## The name
 
 Yes, I'm aware this package should be named `duckdb-driver` or something, I wasn't thinking when I named it and it's too hard to change the name now

duckdb_engine/__init__.py

@@ -296,11 +296,16 @@ def connect(self, *cargs: Any, **cparams: Any) -> "Connection":
                 user_agent = f"{user_agent} {config['custom_user_agent']}"
             config["custom_user_agent"] = user_agent
 
+        filesystems = cparams.pop("register_filesystems", [])
+
         conn = duckdb.connect(*cargs, **cparams)
 
         for extension in preload_extensions:
             conn.execute(f"LOAD {extension}")
 
+        for filesystem in filesystems:
+            conn.register_filesystem(filesystem)
+
         apply_config(self, conn, ext)
 
         return ConnectionWrapper(conn)

duckdb_engine/tests/test_basic.py

@@ -7,6 +7,7 @@
 from typing import Any, Generic, Optional, TypeVar, cast
 
 import duckdb
+import fsspec
 import sqlalchemy
 from hypothesis import assume, given, settings
 from hypothesis.strategies import text as text_strat
@@ -689,3 +690,18 @@ def test_reserved_keywords(engine: Engine) -> None:
     stmt = select(column("qualify"))
 
     assert str(stmt.compile(engine)) == 'SELECT "qualify"'
+
+
+def test_register_filesystem() -> None:
+    memory_fs = fsspec.filesystem("memory")
+    file_fs = fsspec.filesystem("file")
+    engine = create_engine(
+        "duckdb:///",
+        connect_args={
+            "register_filesystems": [memory_fs, file_fs],
+        },
+    )
+
+    with engine.connect() as conn:
+        duckdb_conn = getattr(conn.connection.dbapi_connection, "_ConnectionWrapper__c")
+        assert duckdb.list_filesystems(connection=duckdb_conn) == ["memory", "file"]

poetry.lock

@@ -33,6 +33,7 @@ pytest-cov = {extras = ["coverage"], version = "^5.0.0"}
 pytest-remotedata = "^0.4.0"
 pytest-snapshot = ">=0.9.0,<1"
 toml = "^0.10.2"
+fsspec = "^2025.2.0"
 
 [tool.poetry.group.devtools.dependencies]
 pdbpp = "^0.10.3"