From 908fcda291793169b405cf28c2bbb967bf64ab62 Mon Sep 17 00:00:00 2001 From: joefrost01 Date: Sat, 11 Apr 2026 21:31:29 +0100 Subject: [PATCH] feat: 19 - cloud storage path and credential support --- src/engine.rs | 4 ++++ src/file_resolution.rs | 10 +++++++++- src/fingerprint.rs | 5 ++++- src/query_pipeline.rs | 27 ++++++++++++++++++++++++++- src/reference_tables.rs | 5 ++++- 5 files changed, 47 insertions(+), 4 deletions(-) diff --git a/src/engine.rs b/src/engine.rs index 0d75f1e..52ac01a 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -30,6 +30,7 @@ pub enum CompressionCodec { #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct CloudSettings { pub s3_region: Option, + pub s3_profile: Option, pub s3_access_key_id: Option, pub gcs_project_id: Option, pub azure_storage_account_name: Option, @@ -385,6 +386,9 @@ impl DuckDbEngine { if let Some(region) = &cloud.s3_region { self.set_option("s3_region", region)?; } + if let Some(profile) = &cloud.s3_profile { + self.set_option("s3_profile", profile)?; + } if let Some(access_key_id) = &cloud.s3_access_key_id { self.set_option("s3_access_key_id", access_key_id)?; } diff --git a/src/file_resolution.rs b/src/file_resolution.rs index 22db567..0ac9be9 100644 --- a/src/file_resolution.rs +++ b/src/file_resolution.rs @@ -50,6 +50,7 @@ pub struct FileResolverConfig { pub delimiter: char, pub on_error: OnErrorMode, pub s3_region: Option, + pub s3_profile: Option, pub gcs_project: Option, pub azure_account: Option, } @@ -66,6 +67,7 @@ impl Default for FileResolverConfig { delimiter: ',', on_error: OnErrorMode::Fail, s3_region: None, + s3_profile: None, gcs_project: None, azure_account: None, } @@ -176,6 +178,9 @@ impl FileResolver { if let Some(region) = &self.config.s3_region { set_sql_option(&conn, "s3_region", region)?; } + if let Some(profile) = &self.config.s3_profile { + set_sql_option(&conn, "s3_profile", profile)?; + } if let Some(project) = &self.config.gcs_project { set_sql_option(&conn, "gcs_project_id", project)?; } @@ -383,7 +388,10 @@ fn set_sql_option(connection: &Connection, key: &str, value: &str) -> Result<(), } fn is_cloud_path(path: &str) -> bool { - path.starts_with("s3://") || path.starts_with("gs://") || path.starts_with("az://") + path.starts_with("s3://") + || path.starts_with("gs://") + || path.starts_with("az://") + || path.starts_with("abfss://") } fn escape_sql_literal(value: &str) -> String { diff --git a/src/fingerprint.rs b/src/fingerprint.rs index aa4b0f5..43a0627 100644 --- a/src/fingerprint.rs +++ b/src/fingerprint.rs @@ -45,7 +45,10 @@ pub fn fingerprint_file(path: &Path) -> Result { fn is_cloud_path(path: &Path) -> bool { let value = path.to_string_lossy(); - value.starts_with("s3://") || value.starts_with("gs://") || value.starts_with("az://") + value.starts_with("s3://") + || value.starts_with("gs://") + || value.starts_with("az://") + || value.starts_with("abfss://") } fn download_cloud_blob(path: &Path) -> Result { diff --git a/src/query_pipeline.rs b/src/query_pipeline.rs index ae7d6d8..f9574e4 100644 --- a/src/query_pipeline.rs +++ b/src/query_pipeline.rs @@ -164,11 +164,12 @@ impl QueryPipeline { let engine = DuckDbEngine::new(EngineConfig { cloud: CloudSettings { s3_region: args.s3_region.clone(), + s3_profile: args.s3_profile.clone(), s3_access_key_id: None, gcs_project_id: args.gcs_project.clone(), azure_storage_account_name: args.azure_account.clone(), }, - load_extensions: false, + load_extensions: requires_cloud_extensions(args, &files), })?; if let Some(schema_path) = &args.schema { @@ -422,6 +423,7 @@ fn resolve_files(args: &QueryArgs) -> Result { delimiter: args.delimiter, on_error: args.on_error, s3_region: args.s3_region.clone(), + s3_profile: args.s3_profile.clone(), gcs_project: args.gcs_project.clone(), azure_account: args.azure_account.clone(), }); @@ -669,6 +671,29 @@ fn describe_input_source(args: &QueryArgs) -> String { "explicit paths".to_string() } +fn requires_cloud_extensions( + args: &QueryArgs, + files: &[crate::file_resolution::ResolvedFile], +) -> bool { + files.iter().any(|file| is_cloud_path(&file.path)) + || args + .refs + .iter() + .filter_map(|entry| entry.split_once('=')) + .any(|(_, path)| is_cloud_path(path.trim())) + || args + .output + .as_ref() + .is_some_and(|path| is_cloud_path(path.to_string_lossy().as_ref())) +} + +fn is_cloud_path(path: &str) -> bool { + path.starts_with("s3://") + || path.starts_with("gs://") + || path.starts_with("az://") + || path.starts_with("abfss://") +} + fn write_manifest_if_requested( args: &QueryArgs, summary: &PipelineResult, diff --git a/src/reference_tables.rs b/src/reference_tables.rs index b0c8fda..6f14bad 100644 --- a/src/reference_tables.rs +++ b/src/reference_tables.rs @@ -104,7 +104,10 @@ fn is_valid_identifier(name: &str) -> bool { } fn is_cloud_path(path: &str) -> bool { - path.starts_with("s3://") || path.starts_with("gs://") || path.starts_with("az://") + path.starts_with("s3://") + || path.starts_with("gs://") + || path.starts_with("az://") + || path.starts_with("abfss://") } fn split_excel_sheet_from_path(path: &str) -> (String, Option) {