Merge pull request #711 from splitgraph/store-options-precedence

Don't override passed object store client options from env

object_store_factory/Cargo.toml

@@ -17,6 +17,8 @@ tracing-subscriber = { workspace = true }
 url = { workspace = true }
 
 [dev-dependencies]
+rstest = "*"
+temp-env = { version = "0.3", features = ["async_closure"] }
 tokio = { workspace = true }
 
 

object_store_factory/src/aws.rs

@@ -224,7 +224,7 @@ pub fn add_amazon_s3_environment_variables(
         if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) {
             if key.starts_with("AWS_") {
                 if let Ok(config_key) = key.to_ascii_lowercase().parse() {
-                    options.insert(config_key, value.to_string());
+                    options.entry(config_key).or_insert(value.to_string());
                 }
             }
         }
@@ -236,10 +236,9 @@ pub fn add_amazon_s3_environment_variables(
             .to_uppercase(),
     ) == Ok("true".to_string())
     {
-        options.insert(
-            AmazonS3ConfigKey::Client(ClientConfigKey::AllowHttp),
-            "true".to_string(),
-        );
+        options
+            .entry(AmazonS3ConfigKey::Client(ClientConfigKey::AllowHttp))
+            .or_insert("true".to_string());
     }
 }
 

object_store_factory/src/google.rs

@@ -107,15 +107,34 @@ pub fn map_options_into_google_config_keys(
 pub fn add_google_cloud_storage_environment_variables(
     options: &mut HashMap<GoogleConfigKey, String>,
 ) {
-    if let Ok(service_account_path) = env::var("SERVICE_ACCOUNT") {
-        options.insert(GoogleConfigKey::ServiceAccount, service_account_path);
+    // Insert the service account {path, key} from env only if neither the path nor the key are
+    // already supplied.
+    //
+    // This is because if the same config was already passed we don't want to override it, and if
+    // the other config was passed the client builder will error out (doesn't support both file and
+    // key).
+    if !options.contains_key(&GoogleConfigKey::ServiceAccount)
+        && !options.contains_key(&GoogleConfigKey::ServiceAccountKey)
+    {
+        if let Ok(service_account_path) = env::var("GOOGLE_SERVICE_ACCOUNT") {
+            options.insert(GoogleConfigKey::ServiceAccount, service_account_path);
+        } else if let Ok(service_account_path) = env::var("GOOGLE_SERVICE_ACCOUNT_KEY") {
+            options.insert(GoogleConfigKey::ServiceAccountKey, service_account_path);
+        }
     }
 
+    // Override all other configs from env vars
     for (os_key, os_value) in env::vars_os() {
         if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) {
             if key.starts_with("GOOGLE_") {
                 if let Ok(config_key) = key.to_ascii_lowercase().parse() {
-                    options.insert(config_key, value.to_string());
+                    if !matches!(
+                        config_key,
+                        GoogleConfigKey::ServiceAccount
+                            | GoogleConfigKey::ServiceAccountKey
+                    ) {
+                        options.entry(config_key).or_insert(value.to_string());
+                    }
                 }
             }
         }

object_store_factory/src/lib.rs

@@ -145,7 +145,7 @@ pub async fn build_object_store_from_opts(
             Ok(store)
         }
         ObjectStoreScheme::Local => {
-            let store = local::LocalConfig {
+            let store = LocalConfig {
                 data_dir: url.path().to_string(),
                 disable_hardlinks: false,
             }
@@ -197,6 +197,7 @@ pub async fn build_storage_location_info_from_opts(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use rstest::rstest;
     use serde_json::json;
 
     #[tokio::test]
@@ -294,4 +295,37 @@ mod tests {
         let result = ObjectStoreConfig::build_from_json(&url, &json_str).await;
         assert!(result.is_err());
     }
+
+    #[rstest]
+    #[tokio::test]
+    async fn test_build_aws_object_store(#[values(true, false)] use_env: bool) {
+        let url = Url::parse("s3://my-bucket").unwrap();
+        let options: HashMap<String, String> = if use_env {
+            HashMap::new()
+        } else {
+            HashMap::from([
+                ("access_key_id".to_string(), "my-key".to_string()),
+                ("secret_access_key".to_string(), "my-secret".to_string()),
+            ])
+        };
+
+        let store = temp_env::async_with_vars(
+            [
+                ("AWS_ACCESS_KEY_ID", Some("env-key")),
+                ("AWS_SECRET_ACCESS_KEY", Some("env-secret")),
+            ],
+            build_object_store_from_opts(&url, options),
+        )
+        .await;
+
+        let debug_output = format!("{store:?}");
+        assert!(debug_output.contains("bucket: \"my-bucket\""));
+        if use_env {
+            assert!(debug_output.contains("key_id: \"env-key\""));
+            assert!(debug_output.contains("secret_key: \"env-secret\""));
+        } else {
+            assert!(debug_output.contains("key_id: \"my-key\""));
+            assert!(debug_output.contains("secret_key: \"my-secret\""));
+        }
+    }
 }