Skip to content

Commit 18fc7ba

Browse files
authored
feat: add parquet to extensions (#108)
1 parent c6b0f30 commit 18fc7ba

File tree

2 files changed

+10
-6
lines changed

2 files changed

+10
-6
lines changed

python/rustac/rustac.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ class DuckdbClient:
2626
Args:
2727
extension_directory: A non-standard extension directory to use.
2828
extensions: A list of extensions to LOAD on client initialization.
29-
install_extensions: Whether to install the spatial and icu extensions on client initialization.
29+
install_extensions: Whether to install the required extensions on client initialization.
3030
use_hive_partitioning: Whether to use hive partitioning for geoparquet queries.
3131
"""
3232

src/duckdb.rs

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ use pyo3_arrow::PyTable;
1313
use stac_duckdb::Client;
1414
use std::{path::PathBuf, sync::Mutex};
1515

16+
const REQUIRED_EXTENSIONS: [&str; 3] = ["spatial", "icu", "parquet"];
17+
1618
#[pyclass(frozen)]
1719
pub struct DuckdbClient(Mutex<Client>);
1820

@@ -34,14 +36,16 @@ impl DuckdbClient {
3436
)?;
3537
}
3638
if install_extensions {
37-
connection.execute("INSTALL spatial", [])?;
38-
connection.execute("INSTALL icu", [])?;
39+
for extension in REQUIRED_EXTENSIONS {
40+
connection.execute(&format!("INSTALL {extension}"), [])?;
41+
}
3942
}
4043
for extension in extensions {
41-
connection.execute(&format!("LOAD '{}'", extension), [])?;
44+
connection.execute(&format!("LOAD '{extension}'"), [])?;
45+
}
46+
for extension in REQUIRED_EXTENSIONS {
47+
connection.execute(&format!("LOAD {extension}"), [])?;
4248
}
43-
connection.execute("LOAD spatial", [])?;
44-
connection.execute("LOAD icu", [])?;
4549
let mut client = Client::from(connection);
4650
client.use_hive_partitioning = use_hive_partitioning;
4751
Ok(DuckdbClient(Mutex::new(client)))

0 commit comments

Comments
 (0)