Skip to content

Commit 1888a15

Browse files
gadomskiceholden
andauthored
refactor: fetch extensions (#681)
@ceholden this is a bit of a simpler implementation. I wasn't finding a `fetch_extensions` setup that felt right to me. I tested locally by deleting my extensions and running just the `no_install` test and it failed 🥳 --------- Co-authored-by: Chris Holden <[email protected]>
1 parent 973414b commit 1888a15

File tree

2 files changed

+68
-11
lines changed

2 files changed

+68
-11
lines changed

crates/duckdb/CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
1313
- `Client.search_to_arrow_table` ([#634](https://github.com/stac-utils/stac-rs/pull/634))
1414
- Conditionally disable parsing the WKB ([#635](https://github.com/stac-utils/stac-rs/pull/635))
1515
- `Client.extensions` ([#665](https://github.com/stac-utils/stac-rs/pull/665))
16+
- `Config.install_extensions` ([#681](https://github.com/stac-utils/stac-rs/pull/681))
1617

1718
### Removed
1819

crates/duckdb/src/lib.rs

+67-11
Original file line numberDiff line numberDiff line change
@@ -79,18 +79,14 @@ pub type Result<T> = std::result::Result<T, Error>;
7979
#[derive(Debug)]
8080
pub struct Client {
8181
connection: Connection,
82+
8283
/// The client's configuration.
8384
pub config: Config,
8485
}
8586

8687
/// Configuration for a client.
87-
#[derive(Debug, Clone, Copy)]
88+
#[derive(Debug, Clone)]
8889
pub struct Config {
89-
/// Whether to enable the s3 credential chain, which allows s3:// url access.
90-
///
91-
/// True by default.
92-
pub use_s3_credential_chain: bool,
93-
9490
/// Whether to enable hive partitioning.
9591
///
9692
/// False by default.
@@ -100,6 +96,25 @@ pub struct Config {
10096
///
10197
/// Disable this to enable geopandas reading, for example.
10298
pub convert_wkb: bool,
99+
100+
/// Whether to enable the S3 credential chain, which allows s3:// url access.
101+
///
102+
/// True by default.
103+
pub use_s3_credential_chain: bool,
104+
105+
/// Whether to enable the Azure credential chain, which allows az:// url access.
106+
///
107+
/// True by default.
108+
pub use_azure_credential_chain: bool,
109+
110+
/// Whether to directly install the httpfs extension.
111+
pub use_httpfs: bool,
112+
113+
/// Whether to install extensions when creating a new connection.
114+
pub install_extensions: bool,
115+
116+
/// Use a custom extension repository.
117+
pub custom_extension_repository: Option<String>,
103118
}
104119

105120
/// A SQL query.
@@ -167,21 +182,45 @@ impl Client {
167182
/// use stac_duckdb::{Client, Config};
168183
///
169184
/// let config = Config {
170-
/// use_s3_credential_chain: true,
171185
/// use_hive_partitioning: true,
172186
/// convert_wkb: true,
187+
/// use_s3_credential_chain: true,
188+
/// use_azure_credential_chain: true,
189+
/// use_httpfs: true,
190+
/// install_extensions: true,
191+
/// custom_extension_repository: None,
173192
/// };
174193
/// let client = Client::with_config(config);
175194
/// ```
176195
pub fn with_config(config: Config) -> Result<Client> {
177196
let connection = Connection::open_in_memory()?;
178-
connection.execute("INSTALL spatial", [])?;
197+
if let Some(ref custom_extension_repository) = config.custom_extension_repository {
198+
connection.execute(
199+
"SET custom_extension_repository = '?'",
200+
[custom_extension_repository],
201+
)?;
202+
}
203+
if config.install_extensions {
204+
connection.execute("INSTALL spatial", [])?;
205+
connection.execute("INSTALL icu", [])?;
206+
}
179207
connection.execute("LOAD spatial", [])?;
180-
connection.execute("INSTALL icu", [])?;
181208
connection.execute("LOAD icu", [])?;
209+
if config.use_httpfs && config.install_extensions {
210+
connection.execute("INSTALL httpfs", [])?;
211+
}
182212
if config.use_s3_credential_chain {
213+
if config.install_extensions {
214+
connection.execute("INSTALL aws", [])?;
215+
}
183216
connection.execute("CREATE SECRET (TYPE S3, PROVIDER CREDENTIAL_CHAIN)", [])?;
184217
}
218+
if config.use_azure_credential_chain {
219+
if config.install_extensions {
220+
connection.execute("INSTALL azure", [])?;
221+
}
222+
connection.execute("CREATE SECRET (TYPE azure, PROVIDER CREDENTIAL_CHAIN)", [])?;
223+
}
185224
Ok(Client { connection, config })
186225
}
187226

@@ -519,15 +558,19 @@ impl Default for Config {
519558
fn default() -> Self {
520559
Config {
521560
use_hive_partitioning: false,
522-
use_s3_credential_chain: true,
523561
convert_wkb: true,
562+
use_s3_credential_chain: true,
563+
use_azure_credential_chain: true,
564+
use_httpfs: true,
565+
install_extensions: true,
566+
custom_extension_repository: None,
524567
}
525568
}
526569
}
527570

528571
#[cfg(test)]
529572
mod tests {
530-
use super::Client;
573+
use super::{Client, Config};
531574
use geo::Geometry;
532575
use rstest::{fixture, rstest};
533576
use stac::{Bbox, Validate};
@@ -542,6 +585,19 @@ mod tests {
542585
Client::new().unwrap()
543586
}
544587

588+
#[test]
589+
fn no_install() {
590+
let _mutex = MUTEX.lock().unwrap();
591+
let config = Config {
592+
install_extensions: false,
593+
..Default::default()
594+
};
595+
let client = Client::with_config(config).unwrap();
596+
client
597+
.search("data/100-sentinel-2-items.parquet", Search::default())
598+
.unwrap();
599+
}
600+
545601
#[rstest]
546602
fn search_all(client: Client) {
547603
let item_collection = client

0 commit comments

Comments
 (0)