Add CLI Migrations (#805)

tembo-cli/Cargo.toml

@@ -82,6 +82,7 @@ test-case = "=2.0.0-rc2"
 clap-markdown = { git = "https://github.com/tembo-io/clap-markdown.git", branch = "main", version = "0.1.3" }
 url = "2.5.0"
 rand = "0.8.5"
+sqlx-cli = "0.7.3"
 
 [target.aarch64-unknown-linux-musl.dependencies]
 openssl = { version = "0.10", features = ["vendored"] }

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

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

tembo-cli/examples/migrations-1/tembo-migrations/migration-test/2024052301_create_users_table.sql

@@ -0,0 +1,5 @@
+CREATE TABLE users (
+    id SERIAL PRIMARY KEY,
+    name TEXT NOT NULL,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+);

tembo-cli/examples/migrations-1/tembo-migrations/migration-test/2024052302_add_email_to_users.sql

@@ -0,0 +1,2 @@
+ALTER TABLE users
+ADD COLUMN email TEXT;

tembo-cli/examples/migrations-1/tembo-migrations/migration-test/2024052303_insert_sample_data.sql

@@ -0,0 +1,4 @@
+INSERT INTO users (name, email)
+VALUES ('Alice', 'alice@example.com'),
+       ('Bob', 'bob@example.com'),
+       ('Charlie', 'charlie@example.com');

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

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

tembo-cli/src/cmd/apply.rs

@@ -13,6 +13,7 @@ use itertools::Itertools;
 use log::info;
 use spinoff::spinners;
 use spinoff::Spinner;
+use sqlx::{migrate::Migrator, PgPool};
 use std::fmt::Write;
 use std::path::PathBuf;
 use std::{
@@ -155,7 +156,7 @@ fn tembo_cloud_apply(
         if instance_setting.stack_file.is_some() {
             error("Stack File is only available for local testing.");
         }
-        let result = tembo_cloud_apply_instance(&env, instance_setting);
+        let result = tembo_cloud_apply_instance(&env, instance_setting, _key.to_string());
 
         match result {
             Ok(i) => i,
@@ -181,10 +182,7 @@ fn docker_apply(
 
     for (_key, instance_setting) in instance_settings.iter_mut() {
         let final_instance_setting = docker_apply_instance(verbose, instance_setting.to_owned())?;
-        final_instance_settings.insert(
-            final_instance_setting.instance_name.clone(),
-            final_instance_setting,
-        );
+        final_instance_settings.insert(_key.to_string(), final_instance_setting);
     }
 
     let rendered_dockercompose: String =
@@ -226,14 +224,29 @@ fn docker_apply(
                 .bold()
         ));
 
-        // If all of the above was successful, we can print the url to user
-        instance_started(
-            &format!(
-                "postgres://postgres:postgres@{}.local.tembo.io:{}",
-                instance_setting.instance_name, port
-            ),
-            "local",
+        let database_url = &format!(
+            "postgres://postgres:postgres@{}.local.tembo.io:{}",
+            instance_setting.instance_name, port
         );
+
+        let migrations_dir = PathBuf::from("tembo-migrations").join(_key);
+
+        // Apply SQLx migrations if the tembo-migrations directory exists
+        if migrations_dir.exists() && migrations_dir.is_dir() {
+            let mut sp = Spinner::new(spinners::Dots, "Applying migrations", spinoff::Color::White);
+            tokio::runtime::Runtime::new()?
+                .block_on(apply_migrations(database_url.to_string(), migrations_dir))?;
+            sp.stop_with_message(&format!(
+                "{} {}",
+                "✓".color(colors::indicator_good()).bold(),
+                format!("Migrations applied for Key {}", _key)
+                    .color(colorful::Color::White)
+                    .bold()
+            ));
+        }
+
+        // If all of the above was successful, we can print the url to user
+        instance_started(database_url, "local");
     }
     Ok(())
 }
@@ -356,6 +369,7 @@ fn process_app_services(
 pub fn tembo_cloud_apply_instance(
     env: &Environment,
     instance_settings: &InstanceSettings,
+    key: String,
 ) -> Result<(), anyhow::Error> {
     let profile = env
         .selected_profile
@@ -412,6 +426,15 @@ pub fn tembo_cloud_apply_instance(
             ));
             clean_console();
             let connection_string = construct_connection_string(conn_info, password);
+
+            let migrations_dir = PathBuf::from("tembo-migrations").join(key);
+            if migrations_dir.exists() && migrations_dir.is_dir() {
+                tokio::runtime::Runtime::new()?.block_on(apply_migrations(
+                    connection_string.to_string(),
+                    migrations_dir,
+                ))?;
+            }
+
             instance_started(
                 &connection_string,
                 &instance_settings.stack_type.clone().unwrap(),
@@ -424,6 +447,24 @@ pub fn tembo_cloud_apply_instance(
     Ok(())
 }
 
+async fn apply_migrations(
+    database_url: String,
+    migrations_dir: PathBuf,
+) -> Result<(), anyhow::Error> {
+    let pool = PgPool::connect(&database_url).await?;
+
+    for entry in fs::read_dir(&migrations_dir)? {
+        let entry = entry?;
+        let path = entry.path();
+        if path.is_file() {
+            let migrator = Migrator::new(migrations_dir.clone()).await?;
+            migrator.run(&pool).await?;
+        }
+    }
+
+    Ok(())
+}
+
 fn get_conn_info_with_creds(
     profile: Profile,
     instance_id: Option<String>,

tembo-cli/tests/integration_tests_docker.rs

@@ -181,6 +181,45 @@ async fn vector() -> Result<(), anyhow::Error> {
     Ok(())
 }
 
+#[tokio::test]
+async fn migrations() -> Result<(), anyhow::Error> {
+    let root_dir = env!("CARGO_MANIFEST_DIR");
+    let test_dir = PathBuf::from(root_dir)
+        .join("examples")
+        .join("migrations-1");
+
+    env::set_current_dir(&test_dir)?;
+
+    // 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();
+
+    // tembo delete
+    let mut cmd = Command::cargo_bin(CARGO_BIN)?;
+    cmd.arg("delete");
+    let _ = cmd.ok();
+
+    // check can't connect
+    assert!(assert_can_connect("migration-test".to_str()).await.is_err());
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn data_warehouse() -> Result<(), anyhow::Error> {
     let instance_name = "data-warehouse";