Fix configuration interpolation (#5403)

* Fix config templating regex * Add config parsing tests to Lambda
quickwit-oss · Sep 9, 2024 · 4ed7995 · 4ed7995
1 parent 79acfe4
commit 4ed7995
Show file tree

Hide file tree

Showing 8 changed files with 126 additions and 2 deletions.
diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock
diff --git a/quickwit/Cargo.toml b/quickwit/Cargo.toml
@@ -215,6 +215,7 @@ serde_json_borrow = "0.5"
 serde_qs = { version = "0.12", features = ["warp"] }
 serde_with = "3.9.0"
 serde_yaml = "0.9"
+serial_test = { version = "3.1.1", features = ["file_locks"] }
 siphasher = "0.3"
 smallvec = "1"
 sqlx = { version = "0.7", features = [

diff --git a/quickwit/quickwit-config/src/templating.rs b/quickwit/quickwit-config/src/templating.rs
@@ -30,7 +30,7 @@ use tracing::debug;
 // `ENV_VAR` or `ENV_VAR:DEFAULT`
 // Ignores whitespaces in curly braces
 static TEMPLATE_ENV_VAR_CAPTURE: Lazy<Regex> = Lazy::new(|| {
-    Regex::new(r"\$\{\s*([A-Za-z0-9_]+)\s*(?::\-\s*([\S]+)\s*)?}")
+    Regex::new(r"\$\{\s*([A-Za-z0-9_]+)\s*(?::\-\s*([^\s\}]+)\s*)?}")
         .expect("regular expression should compile")
 });
 
@@ -158,6 +158,23 @@ mod test {
         assert_eq!(rendered, "metastore_uri: s3://test-bucket/metastore");
     }
 
+    #[test]
+    fn test_template_render_with_multiple_vars_per_line() {
+        let config_content =
+            b"metastore_uri: s3://${RENDER_MULTIPLE_BUCKET}/${RENDER_MULTIPLE_PREFIX:-index}#polling_interval=${RENDER_MULTIPLE_INTERVAL}s";
+        env::set_var("RENDER_MULTIPLE_BUCKET", "test-bucket");
+        env::set_var("RENDER_MULTIPLE_PREFIX", "metastore");
+        env::set_var("RENDER_MULTIPLE_INTERVAL", "30");
+        let rendered = render_config(config_content).unwrap();
+        std::env::remove_var("RENDER_MULTIPLE_BUCKET");
+        std::env::remove_var("RENDER_MULTIPLE_PREFIX");
+        std::env::remove_var("RENDER_MULTIPLE_INTERVAL");
+        assert_eq!(
+            rendered,
+            "metastore_uri: s3://test-bucket/metastore#polling_interval=30s"
+        );
+    }
+
     #[test]
     fn test_template_render_ignores_commented_lines() {
         {

diff --git a/quickwit/quickwit-lambda/Cargo.toml b/quickwit/quickwit-lambda/Cargo.toml
@@ -24,6 +24,7 @@ s3-localstack-tests = []
 [dependencies]
 anyhow = { workspace = true }
 aws_lambda_events = "0.15.0"
+bytesize = { workspace = true }
 chitchat = { workspace = true }
 chrono = { workspace = true }
 flate2 = { workspace = true }
@@ -65,3 +66,6 @@ quickwit-search = { workspace = true }
 quickwit-serve = { workspace = true }
 quickwit-storage = { workspace = true }
 quickwit-telemetry = { workspace = true }
+
+[dev-dependencies]
+serial_test = { workspace = true }
diff --git a/quickwit/quickwit-lambda/src/indexer/environment.rs b/quickwit/quickwit-lambda/src/indexer/environment.rs
@@ -48,3 +48,47 @@ pub static MAX_CHECKPOINTS: Lazy<usize> = Lazy::new(|| {
             .expect("QW_LAMBDA_MAX_CHECKPOINTS must be a positive integer")
     })
 });
+
+#[cfg(test)]
+mod tests {
+
+    use quickwit_config::{ConfigFormat, NodeConfig};
+
+    use super::*;
+
+    #[tokio::test]
+    #[serial_test::file_serial(with_env)]
+    async fn test_load_config() {
+        let bucket = "mock-test-bucket";
+        std::env::set_var("QW_LAMBDA_METASTORE_BUCKET", bucket);
+        std::env::set_var("QW_LAMBDA_INDEX_BUCKET", bucket);
+        std::env::set_var(
+            "QW_LAMBDA_INDEX_CONFIG_URI",
+            "s3://mock-index-config-bucket",
+        );
+        std::env::set_var("QW_LAMBDA_INDEX_ID", "lambda-test");
+
+        let node_config = NodeConfig::load(ConfigFormat::Yaml, CONFIGURATION_TEMPLATE.as_bytes())
+            .await
+            .unwrap();
+        //
+        assert_eq!(
+            node_config.data_dir_path.to_string_lossy(),
+            "/tmp",
+            "only `/tmp` is writeable in AWS Lambda"
+        );
+        assert_eq!(
+            node_config.default_index_root_uri,
+            "s3://mock-test-bucket/index"
+        );
+        assert_eq!(
+            node_config.metastore_uri.to_string(),
+            "s3://mock-test-bucket/index"
+        );
+
+        std::env::remove_var("QW_LAMBDA_METASTORE_BUCKET");
+        std::env::remove_var("QW_LAMBDA_INDEX_BUCKET");
+        std::env::remove_var("QW_LAMBDA_INDEX_CONFIG_URI");
+        std::env::remove_var("QW_LAMBDA_INDEX_ID");
+    }
+}
diff --git a/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs b/quickwit/quickwit-lambda/src/indexer/ingest/mod.rs
@@ -150,6 +150,7 @@ mod tests {
     }
 
     #[tokio::test]
+    #[serial_test::file_serial(with_env)]
     async fn test_ingest() -> anyhow::Result<()> {
         quickwit_common::setup_logging_for_tests();
         let bucket = "quickwit-integration-tests";
@@ -246,6 +247,13 @@ mod tests {
             assert_eq!(stats.num_docs, 1);
         }
 
+        std::env::remove_var("QW_LAMBDA_METASTORE_BUCKET");
+        std::env::remove_var("QW_LAMBDA_INDEX_BUCKET");
+        std::env::remove_var("QW_LAMBDA_METASTORE_PREFIX");
+        std::env::remove_var("QW_LAMBDA_INDEX_PREFIX");
+        std::env::remove_var("QW_LAMBDA_INDEX_CONFIG_URI");
+        std::env::remove_var("QW_LAMBDA_INDEX_ID");
+
         Ok(())
     }
 }
diff --git a/quickwit/quickwit-lambda/src/searcher/environment.rs b/quickwit/quickwit-lambda/src/searcher/environment.rs
@@ -26,3 +26,51 @@ data_dir: /tmp
 searcher:
   partial_request_cache_capacity: ${QW_LAMBDA_PARTIAL_REQUEST_CACHE_CAPACITY:-64M}
 "#;
+
+#[cfg(test)]
+mod tests {
+
+    use bytesize::ByteSize;
+    use quickwit_config::{ConfigFormat, NodeConfig};
+
+    use super::*;
+
+    #[tokio::test]
+    #[serial_test::file_serial(with_env)]
+    async fn test_load_config() {
+        let bucket = "mock-test-bucket";
+        std::env::set_var("QW_LAMBDA_METASTORE_BUCKET", bucket);
+        std::env::set_var("QW_LAMBDA_INDEX_BUCKET", bucket);
+        std::env::set_var(
+            "QW_LAMBDA_INDEX_CONFIG_URI",
+            "s3://mock-index-config-bucket",
+        );
+        std::env::set_var("QW_LAMBDA_INDEX_ID", "lambda-test");
+
+        let node_config = NodeConfig::load(ConfigFormat::Yaml, CONFIGURATION_TEMPLATE.as_bytes())
+            .await
+            .unwrap();
+        assert_eq!(
+            node_config.data_dir_path.to_string_lossy(),
+            "/tmp",
+            "only `/tmp` is writeable in AWS Lambda"
+        );
+        assert_eq!(
+            node_config.default_index_root_uri,
+            "s3://mock-test-bucket/index"
+        );
+        assert_eq!(
+            node_config.metastore_uri.to_string(),
+            "s3://mock-test-bucket/index#polling_interval=60s"
+        );
+        assert_eq!(
+            node_config.searcher_config.partial_request_cache_capacity,
+            ByteSize::mb(64)
+        );
+
+        std::env::remove_var("QW_LAMBDA_METASTORE_BUCKET");
+        std::env::remove_var("QW_LAMBDA_INDEX_BUCKET");
+        std::env::remove_var("QW_LAMBDA_INDEX_CONFIG_URI");
+        std::env::remove_var("QW_LAMBDA_INDEX_ID");
+    }
+}
diff --git a/quickwit/quickwit-metastore/Cargo.toml b/quickwit/quickwit-metastore/Cargo.toml
@@ -52,7 +52,7 @@ futures = { workspace = true }
 md5 = { workspace = true }
 mockall = { workspace = true }
 rand = { workspace = true }
-serial_test = { version = "3.1.1", features = ["file_locks"] }
+serial_test = { workspace = true }
 tempfile = { workspace = true }
 tracing-subscriber = { workspace = true }