From ddc743be314d214bf64cd000b8010f9292aa9880 Mon Sep 17 00:00:00 2001 From: Pete Gadomski Date: Wed, 23 Apr 2025 07:53:23 -0600 Subject: [PATCH] feat: add parquet to extensions --- python/rustac/rustac.pyi | 2 +- src/duckdb.rs | 14 +++++++++----- 2 files changed, 10 insertions(+), 6 deletions(-) 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)))