refactor: simplify extension handling in duckdb

Author: gadomski

.gitignore

@@ -164,3 +164,5 @@ cython_debug/
 #.idea/
 .vscode/
 docs/generated
+
+tests/duckdb-extensions

Cargo.lock

@@ -13,6 +13,7 @@ duckdb-bundled = ["stac-duckdb/bundled"]
 
 [dependencies]
 clap = "4.5.31"
+duckdb = { version = "1.2.2", features = ["serde_json"] }
 geoarrow-array = { git = "https://github.com/geoarrow/geoarrow-rs/", rev = "17bf33e4cf78b060afa08ca9560dc4efd73c2c76" }
 geojson = "0.24.1"
 pyo3 = { version = "0.24.1", features = ["extension-module"] }

Cargo.toml

@@ -1,5 +1,6 @@
 """The power of Rust for the Python STAC ecosystem."""
 
+from pathlib import Path
 from typing import Any, AsyncIterator, Literal, Optional, Tuple
 
 import arro3.core
@@ -15,28 +16,28 @@ class DuckdbClient:
     def __init__(
         self,
         *,
-        use_s3_credential_chain: bool = True,
-        use_azure_credential_chain: bool = True,
-        use_httpfs: bool = True,
+        extension_directory: Path | None = None,
+        extensions: list[str] | None = None,
+        install_spatial: bool = True,
         use_hive_partitioning: bool = False,
-        install_extensions: bool = True,
-        custom_extension_repository: str | None = None,
-        extension_directory: str | None = None,
     ) -> None:
         """Creates a new duckdb client.
 
         Args:
-            use_s3_credential_chain: If true, configures DuckDB to correctly
-                handle s3:// urls.
-            use_azure_credential_chain: If true, configures DuckDB to correctly
-                handle azure urls.
-            use_httpfs: If true, configures DuckDB to correctly handle https
-                urls.
-            use_hive_partitioning: If true, enables queries on hive partitioned
-                geoparquet files.
-            install_extensions: If true, installs extensions before loading them.
-            custom_extension_repository: A custom extension repository to use.
             extension_directory: A non-standard extension directory to use.
+            extensions: A list of extensions to LOAD on client initialization.
+            install_spatial: Whether to install the spatial extension on client initialization.
+            use_hive_partitioning: Whether to use hive partitioning for geoparquet queries.
+        """
+
+    def execute(self, sql: str, params: list[str] | None = None) -> int:
+        """Execute an SQL command.
+
+        This can be useful for configuring AWS credentials, for example.
+
+        Args:
+            sql: The SQL to execute
+            params: The parameters to pass in to the execution
         """
 
     def search(

python/rustac/rustac.pyi

@@ -2,6 +2,5 @@
 
 set -e
 
-uv run maturin dev --uv -E arrow
 uv run pytest "$@"
 uv run rustac translate spec-examples/v1.1.0/simple-item.json /dev/null

scripts/test

@@ -2,6 +2,7 @@ use crate::{
     search::{PySortby, StringOrDict, StringOrList},
     Result,
 };
+use duckdb::Connection;
 use pyo3::{
     exceptions::PyException,
     prelude::*,
@@ -18,13 +19,40 @@ pub struct DuckdbClient(Mutex<Client>);
 #[pymethods]
 impl DuckdbClient {
     #[new]
-    #[pyo3(signature = (*, use_hive_partitioning=false, extension_directory=None, install_extensions=true))]
+    #[pyo3(signature = (*, extension_directory=None, extensions=Vec::new(), install_spatial=true, use_hive_partitioning=false))]
     fn new(
-        use_hive_partitioning: bool,
         extension_directory: Option<PathBuf>,
-        install_extensions: bool,
+        extensions: Vec<String>,
+        install_spatial: bool,
+        use_hive_partitioning: bool,
     ) -> Result<DuckdbClient> {
-        todo!()
+        let connection = Connection::open_in_memory()?;
+        if let Some(extension_directory) = extension_directory {
+            connection.execute(
+                "SET extension_directory = ?",
+                [extension_directory.to_string_lossy()],
+            )?;
+        }
+        if install_spatial {
+            connection.execute("INSTALL spatial", [])?;
+        }
+        for extension in extensions {
+            connection.execute(&format!("LOAD '{}'", extension), [])?;
+        }
+        connection.execute("LOAD spatial", [])?;
+        let mut client = Client::from(connection);
+        client.use_hive_partitioning = use_hive_partitioning;
+        Ok(DuckdbClient(Mutex::new(client)))
+    }
+
+    #[pyo3(signature = (sql, params = Vec::new()))]
+    fn execute<'py>(&self, sql: String, params: Vec<String>) -> Result<usize> {
+        let client = self
+            .0
+            .lock()
+            .map_err(|err| PyException::new_err(err.to_string()))?;
+        let count = client.execute(&sql, duckdb::params_from_iter(params))?;
+        Ok(count)
     }
 
     #[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, **kwargs))]

src/duckdb.rs

@@ -9,6 +9,9 @@ create_exception!(rustac, RustacError, PyException);
 
 #[derive(Debug, Error)]
 pub enum Error {
+    #[error(transparent)]
+    Duckdb(#[from] duckdb::Error),
+
     #[error(transparent)]
     Geojson(#[from] geojson::Error),
 

src/error.rs

@@ -1,9 +1,8 @@
 from pathlib import Path
 
 import pytest
-from geopandas import GeoDataFrame
-
 import rustac
+from geopandas import GeoDataFrame
 from rustac import DuckdbClient, RustacError
 
 
@@ -12,6 +11,11 @@ def client() -> DuckdbClient:
     return DuckdbClient()
 
 
+@pytest.fixture
+def extension_directory() -> Path:
+    return Path(__file__).parent / "duckdb-extensions"
+
+
 def test_search(client: DuckdbClient) -> None:
     items = client.search("data/extended-item.parquet")
     assert len(items) == 1
@@ -40,13 +44,37 @@ def test_search_to_arrow(client: DuckdbClient) -> None:
     assert len(item_collection["features"]) == 100
 
 
-def test_custom_extension_directory() -> None:
-    extension_directory = Path(__file__).parent / "duckdb-extensions"
+def test_custom_extension_directory(extension_directory: Path) -> None:
     client = DuckdbClient(extension_directory=extension_directory)
     # Search to ensure we trigger everything
     client.search("data/100-sentinel-2-items.parquet")
 
 
 def test_no_install(tmp_path: Path) -> None:
     with pytest.raises(RustacError):
-        DuckdbClient(extension_directory=tmp_path, install_extensions=False)
+        DuckdbClient(extension_directory=tmp_path, install_spatial=False)
+
+
+def test_extensions(extension_directory: Path, tmp_path: Path) -> None:
+    # Ensure we've fetched the extension
+    DuckdbClient(extension_directory=extension_directory)
+
+    extension = next(extension_directory.glob("**/spatial.duckdb_extension"))
+    client = DuckdbClient(
+        extensions=[str(extension)], extension_directory=tmp_path, install_spatial=False
+    )
+    client.search("data/100-sentinel-2-items.parquet")
+
+
+def test_execute(client: DuckdbClient, extension_directory: Path) -> None:
+    # Just a smoke test
+    client.execute("SET extension_directory = ?", [str(extension_directory)])
+
+
+def test_load_spatial() -> None:
+    DuckdbClient(extensions=["spatial"])
+
+
+@pytest.mark.skip("slow")
+def test_aws_credential_chain(client: DuckdbClient) -> None:
+    client.execute("CREATE SECRET (TYPE S3, PROVIDER CREDENTIAL_CHAIN)")