diff --git a/python/rustac/rustac.pyi b/python/rustac/rustac.pyi index 69ee246..b194f72 100644 --- a/python/rustac/rustac.pyi +++ b/python/rustac/rustac.pyi @@ -26,7 +26,7 @@ class DuckdbClient: Args: extension_directory: A non-standard extension directory to use. extensions: A list of extensions to LOAD on client initialization. - install_extensions: Whether to install the spatial and icu extensions on client initialization. + install_extensions: Whether to install the required extensions on client initialization. use_hive_partitioning: Whether to use hive partitioning for geoparquet queries. """ diff --git a/src/duckdb.rs b/src/duckdb.rs index 1e77425..37ddec5 100644 --- a/src/duckdb.rs +++ b/src/duckdb.rs @@ -13,6 +13,8 @@ use pyo3_arrow::PyTable; use stac_duckdb::Client; use std::{path::PathBuf, sync::Mutex}; +const REQUIRED_EXTENSIONS: [&str; 3] = ["spatial", "icu", "parquet"]; + #[pyclass(frozen)] pub struct DuckdbClient(Mutex); @@ -34,14 +36,16 @@ impl DuckdbClient { )?; } if install_extensions { - connection.execute("INSTALL spatial", [])?; - connection.execute("INSTALL icu", [])?; + for extension in REQUIRED_EXTENSIONS { + connection.execute(&format!("INSTALL {extension}"), [])?; + } } for extension in extensions { - connection.execute(&format!("LOAD '{}'", extension), [])?; + connection.execute(&format!("LOAD '{extension}'"), [])?; + } + for extension in REQUIRED_EXTENSIONS { + connection.execute(&format!("LOAD {extension}"), [])?; } - connection.execute("LOAD spatial", [])?; - connection.execute("LOAD icu", [])?; let mut client = Client::from(connection); client.use_hive_partitioning = use_hive_partitioning; Ok(DuckdbClient(Mutex::new(client)))