Allow Secrets for Migrations - CLI (#831)

tembo-cli/Cargo.toml

@@ -100,3 +100,8 @@ openssl = { version = "0.10", features = ["vendored"] }
 assert_cmd = "2.0.8"
 predicates = "2.1.5"
 rstest = "0.18"
+
+# We can't run integration tests in parallel currently since there are
+# conflicts in containers (like traefik) provisioned by various tests
+[env]
+RUST_TEST_THREADS = "1"

tembo-cli/examples/migrations-2/.gitignore

@@ -0,0 +1,2 @@
+docker-compose.yml
+migration-test-2/*

tembo-cli/examples/migrations-2/tembo-migrations/migration-test-2/20240621_create_custom_table.sql

@@ -0,0 +1,16 @@
+CREATE TABLE IF NOT EXISTS custom_secret_table (
+    id SERIAL PRIMARY KEY,
+    secret_value TEXT NOT NULL
+);
+
+CREATE OR REPLACE FUNCTION set_custom_parameters() RETURNS void AS $$
+BEGIN
+    PERFORM set_config('tembo.custom_secret', current_setting('tembo.custom_secret'), false);
+END;
+$$ LANGUAGE plpgsql;
+
+SELECT set_custom_parameters();
+
+INSERT INTO custom_secret_table (secret_value) VALUES (current_setting('tembo.custom_secret'));
+
+SELECT secret_value FROM custom_secret_table WHERE secret_value = current_setting('tembo.custom_secret');

tembo-cli/examples/migrations-2/tembo.toml

@@ -0,0 +1,4 @@
+[migration-test-2]
+environment = "dev"
+instance_name = "migration-test-2"
+stack_type = "Standard"

tembo-cli/src/cmd/apply.rs

@@ -13,7 +13,9 @@ use itertools::Itertools;
 use log::info;
 use spinoff::spinners;
 use spinoff::Spinner;
-use sqlx::{migrate::Migrator, PgPool};
+use sqlx::migrate::Migrator;
+use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
+use std::env;
 use std::fmt::Write;
 use std::path::PathBuf;
 use std::{
@@ -461,7 +463,17 @@ async fn apply_migrations(
     database_url: String,
     migrations_dir: PathBuf,
 ) -> Result<(), anyhow::Error> {
-    let pool = PgPool::connect(&database_url).await?;
+    // Create connection options
+    let mut options: PgConnectOptions = database_url.parse()?;
+
+    for (key, value) in env::vars() {
+        if key.starts_with("TEMBO_") {
+            let parameter = key.replace("TEMBO_", "tembo.").to_lowercase();
+            options = options.options([(parameter.as_str(), value.as_str())]);
+        }
+    }
+
+    let pool = PgPoolOptions::new().connect_with(options).await?;
 
     for entry in fs::read_dir(&migrations_dir)? {
         let entry = entry?;
@@ -1480,6 +1492,11 @@ mod tests {
         cmd.arg("delete");
         let _ = cmd.ok();
 
+        let mut cmd = Command::new("sh");
+        cmd.arg("-c");
+        cmd.arg("docker volume rm $(docker volume ls -q)");
+        cmd.assert().success();
+
         Ok(())
     }
 }

tembo-cli/src/cmd/logs.rs

@@ -415,6 +415,11 @@ mod tests {
         cmd.arg("delete");
         let _ = cmd.ok();
 
+        let mut cmd = Command::new("sh");
+        cmd.arg("-c");
+        cmd.arg("docker volume rm $(docker volume ls -q)");
+        cmd.assert().success();
+
         Ok(())
     }
 

tembo-cli/tests/integration_tests_docker.rs

@@ -175,6 +175,11 @@ async fn vector() -> Result<(), anyhow::Error> {
     cmd.arg("delete");
     let _ = cmd.ok();
 
+    let mut cmd = Command::new("sh");
+    cmd.arg("-c");
+    cmd.arg("docker volume rm $(docker volume ls -q)");
+    cmd.assert().success();
+
     // check can't connect
     assert!(assert_can_connect("vector".to_str()).await.is_err());
 
@@ -214,6 +219,11 @@ async fn migrations() -> Result<(), anyhow::Error> {
     cmd.arg("delete");
     let _ = cmd.ok();
 
+    let mut cmd = Command::new("sh");
+    cmd.arg("-c");
+    cmd.arg("docker volume rm $(docker volume ls -q)");
+    cmd.assert().success();
+
     // check can't connect
     assert!(assert_can_connect("migration-test".to_str()).await.is_err());
 
@@ -265,6 +275,11 @@ async fn data_warehouse() -> Result<(), anyhow::Error> {
     cmd.arg("delete");
     let _ = cmd.ok();
 
+    let mut cmd = Command::new("sh");
+    cmd.arg("-c");
+    cmd.arg("docker volume rm $(docker volume ls -q)");
+    cmd.assert().success();
+
     // check can't connect
     assert!(assert_can_connect(instance_name.to_string()).await.is_err());
 
@@ -341,6 +356,11 @@ async fn multiple_instances() -> Result<(), anyhow::Error> {
     cmd.arg("delete");
     let _ = cmd.ok();
 
+    let mut cmd = Command::new("sh");
+    cmd.arg("-c");
+    cmd.arg("docker volume rm $(docker volume ls -q)");
+    cmd.assert().success();
+
     // check can't connect
     assert!(assert_can_connect(instance1_name.to_str()).await.is_err());
     assert!(assert_can_connect(instance2_name.to_str()).await.is_err());
@@ -410,13 +430,49 @@ async fn local_persistence() -> Result<(), anyhow::Error> {
     cmd.arg("delete");
     let _ = cmd.ok();
 
-    // Remove the Docker volume
-    let mut cmd = Command::new("docker");
-    cmd.arg("volume");
-    cmd.arg("rm");
-    cmd.arg("set_set-data");
+    let mut cmd = Command::new("sh");
+    cmd.arg("-c");
+    cmd.arg("docker volume rm $(docker volume ls -q)");
+    cmd.assert().success();
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn run_migration_secret() -> Result<(), anyhow::Error> {
+    let instance_name = "migrations-2";
+    let root_dir = env!("CARGO_MANIFEST_DIR");
+    let test_dir = PathBuf::from(root_dir).join("examples").join(instance_name);
+
+    env::set_current_dir(&test_dir)?;
+
+    // Set the environment variable
+    env::set_var("TEMBO_CUSTOM_SECRET", "my_custom_secret_value");
+
+    // tembo init
+    let mut cmd = Command::cargo_bin(CARGO_BIN)?;
+    cmd.arg("init");
+    cmd.assert().success();
+
+    // tembo context set --name local
+    let mut cmd = Command::cargo_bin(CARGO_BIN)?;
+    cmd.arg("context");
+    cmd.arg("set");
+    cmd.arg("--name");
+    cmd.arg("local");
     cmd.assert().success();
 
+    // tembo apply
+    let mut cmd = Command::cargo_bin(CARGO_BIN)?;
+    cmd.arg("--verbose");
+    cmd.arg("apply");
+    cmd.assert().success();
+
+    // Stop the container
+    let mut cmd = Command::cargo_bin(CARGO_BIN)?;
+    cmd.arg("delete");
+    let _ = cmd.ok();
+
     Ok(())
 }