datafusion-ducklake-cli 0.1.0

Interactive and scriptable SQL shell for datafusion-ducklake-provider.
Documentation
use std::collections::BTreeMap;
use std::env;

use datafusion_ducklake_provider::DuckLakeObjectStoreConfig;

use crate::args::Cli;
use crate::error::{CliError, err};

pub(crate) fn object_store_configs(cli: &Cli) -> Result<Vec<DuckLakeObjectStoreConfig>, CliError> {
    if cli.object_store_url.is_empty() {
        if cli.object_store_local_path.is_some()
            || !cli.object_store_config.is_empty()
            || cli_has_explicit_s3_options(cli)
        {
            return Err(err(
                "object-store options require at least one --object-store-url",
            ));
        }
        return Ok(Vec::new());
    }

    if cli.object_store_local_path.is_some() && cli.object_store_url.len() != 1 {
        return Err(err(
            "--object-store-local-path can only be used with one --object-store-url",
        ));
    }

    let options = object_store_options(cli)?;
    cli.object_store_url
        .iter()
        .map(|url| {
            validate_object_store_url(url)?;
            let mut config = DuckLakeObjectStoreConfig::new(url.clone());
            config.object_store_local_path = cli
                .object_store_local_path
                .as_ref()
                .map(|path| path.display().to_string());
            config.object_store_config = options.clone();
            Ok(config)
        })
        .collect()
}

fn object_store_options(cli: &Cli) -> Result<Vec<(String, String)>, CliError> {
    let mut options = BTreeMap::new();

    for raw in &cli.object_store_config {
        let Some((key, value)) = raw.split_once('=') else {
            return Err(err(format!(
                "--object-store-config expects KEY=VALUE, got {raw:?}"
            )));
        };
        set_option(&mut options, normalize_config_key(key), value.to_string());
    }

    set_option_from_env_if_missing(&mut options, "aws_endpoint", "DUCKLAKE_S3_ENDPOINT");
    set_option_from_env_if_missing(
        &mut options,
        "aws_access_key_id",
        "DUCKLAKE_S3_ACCESS_KEY_ID",
    );
    set_option_from_env_if_missing(
        &mut options,
        "aws_secret_access_key",
        "DUCKLAKE_S3_SECRET_ACCESS_KEY",
    );
    set_option_from_env_if_missing(
        &mut options,
        "aws_session_token",
        "DUCKLAKE_S3_SESSION_TOKEN",
    );
    set_option_from_env_if_missing(&mut options, "aws_region", "DUCKLAKE_S3_REGION");

    if let Some(value) = env_bool("DUCKLAKE_S3_ALLOW_HTTP")? {
        options
            .entry("aws_allow_http".to_string())
            .or_insert(value.to_string());
    }
    if env_bool("DUCKLAKE_S3_PATH_STYLE")?.unwrap_or(false) {
        options
            .entry("aws_virtual_hosted_style_request".to_string())
            .or_insert("false".to_string());
    }

    set_option_if_some(&mut options, "aws_endpoint", cli.s3_endpoint.clone());
    set_option_if_some(
        &mut options,
        "aws_access_key_id",
        cli.s3_access_key_id.clone(),
    );
    set_option_if_some(
        &mut options,
        "aws_secret_access_key",
        cli.s3_secret_access_key.clone(),
    );
    set_option_if_some(
        &mut options,
        "aws_session_token",
        cli.s3_session_token.clone(),
    );
    set_option_if_some(&mut options, "aws_region", cli.s3_region.clone());
    set_option_if_some(
        &mut options,
        "aws_allow_http",
        cli.s3_allow_http.map(|value| value.to_string()),
    );
    if cli.s3_path_style {
        set_option(
            &mut options,
            "aws_virtual_hosted_style_request",
            "false".to_string(),
        );
    }

    Ok(options.into_iter().collect())
}

fn validate_object_store_url(url: &str) -> Result<(), CliError> {
    if url.starts_with("s3://") && !cfg!(feature = "object-store-s3") {
        return Err(err(
            "s3:// object stores require ducklake-cli to be built with the object-store-s3 feature",
        ));
    }
    Ok(())
}

fn cli_has_explicit_s3_options(cli: &Cli) -> bool {
    cli.s3_endpoint.is_some()
        || cli.s3_access_key_id.is_some()
        || cli.s3_secret_access_key.is_some()
        || cli.s3_session_token.is_some()
        || cli.s3_region.is_some()
        || cli.s3_allow_http.is_some()
        || cli.s3_path_style
}

fn normalize_config_key(key: &str) -> String {
    key.trim()
        .strip_prefix("object_store_config.")
        .or_else(|| key.trim().strip_prefix("object_store_config_"))
        .unwrap_or_else(|| key.trim())
        .to_string()
}

fn set_option(options: &mut BTreeMap<String, String>, key: impl Into<String>, value: String) {
    options.insert(key.into(), value);
}

fn set_option_if_some(options: &mut BTreeMap<String, String>, key: &str, value: Option<String>) {
    if let Some(value) = value {
        set_option(options, key, value);
    }
}

fn set_option_from_env_if_missing(
    options: &mut BTreeMap<String, String>,
    key: &str,
    env_key: &str,
) {
    if options.contains_key(key) {
        return;
    }
    if let Ok(value) = env::var(env_key) {
        set_option(options, key, value);
    }
}

fn env_bool(key: &str) -> Result<Option<bool>, CliError> {
    match env::var(key) {
        Ok(value) => parse_bool(&value).map(Some),
        Err(env::VarError::NotPresent) => Ok(None),
        Err(env::VarError::NotUnicode(_)) => Err(err(format!("{key} is not valid UTF-8"))),
    }
}

fn parse_bool(value: &str) -> Result<bool, CliError> {
    match value.trim().to_ascii_lowercase().as_str() {
        "1" | "true" | "t" | "yes" | "y" | "on" => Ok(true),
        "0" | "false" | "f" | "no" | "n" | "off" => Ok(false),
        _ => Err(err(format!("expected boolean value, got {value:?}"))),
    }
}

#[cfg(test)]
mod tests {
    use clap::Parser;

    use super::{object_store_configs, parse_bool};
    use crate::Cli;

    #[test]
    fn maps_s3_flags_to_object_store_options() {
        let cli = Cli::parse_from([
            "ducklake-cli",
            "--object-store-url",
            "https://example.com",
            "--object-store-config",
            "aws_region=eu-central-1",
            "--s3-endpoint",
            "http://localhost:9000",
            "--s3-access-key-id",
            "admin",
            "--s3-secret-access-key",
            "password",
            "--s3-allow-http",
            "true",
            "--s3-path-style",
        ]);

        let configs = object_store_configs(&cli).unwrap();
        assert_eq!(configs.len(), 1);
        let mut options = configs[0].object_store_config.clone();
        options.sort();
        assert_eq!(
            options,
            vec![
                ("aws_access_key_id".to_string(), "admin".to_string()),
                ("aws_allow_http".to_string(), "true".to_string()),
                (
                    "aws_endpoint".to_string(),
                    "http://localhost:9000".to_string()
                ),
                ("aws_region".to_string(), "eu-central-1".to_string()),
                ("aws_secret_access_key".to_string(), "password".to_string()),
                (
                    "aws_virtual_hosted_style_request".to_string(),
                    "false".to_string()
                ),
            ]
        );
    }

    #[test]
    fn object_store_options_require_url() {
        let cli = Cli::parse_from(["ducklake-cli", "--s3-endpoint", "http://localhost:9000"]);
        let err = object_store_configs(&cli).unwrap_err().to_string();
        assert!(err.contains("--object-store-url"));
    }

    #[test]
    fn parses_boolish_values() {
        assert!(parse_bool("true").unwrap());
        assert!(parse_bool("1").unwrap());
        assert!(!parse_bool("false").unwrap());
        assert!(!parse_bool("off").unwrap());
        assert!(parse_bool("wat").is_err());
    }
}