runs sqlx migration for tembo cloud deployment (#448)

tembo-cli/Cargo.toml

@@ -53,6 +53,7 @@ tokio = { version = "1.26.0", features = [
 ] }
 controller = "0.26.1"
 sqlx = { version = "0.7.3", features = ["runtime-tokio-native-tls", "postgres", "chrono", "json"] }
+base64 = "0.21.5"
 
 [dev-dependencies]
 assert_cmd = "2.0.8"

tembo-cli/src/cli/docker.rs

@@ -2,10 +2,6 @@ use crate::Result;
 use anyhow::bail;
 use simplelog::*;
 use spinners::{Spinner, Spinners};
-use sqlx::migrate::Migrator;
-use sqlx::Pool;
-use sqlx::Postgres;
-use std::path::Path;
 use std::process::Command as ShellCommand;
 use std::process::Output;
 
@@ -65,20 +61,6 @@ impl Docker {
         Ok(())
     }
 
-    // run sqlx migrate
-    pub async fn run_sqlx_migrate() -> Result {
-        let mut sp = Spinner::new(Spinners::Line, "Running SQL migration".into());
-
-        let pool = Pool::<Postgres>::connect("postgres://postgres:postgres@localhost:5432").await?;
-
-        let m = Migrator::new(Path::new("./migrations")).await?;
-        m.run(&pool).await?;
-
-        sp.stop_with_message("- SQL migration completed".to_string());
-
-        Ok(())
-    }
-
     // stop & remove container for given name
     pub fn stop_remove(name: &str) -> Result {
         let mut sp = Spinner::new(Spinners::Line, "Stopping & Removing instance".into());

tembo-cli/src/cli/mod.rs

@@ -2,4 +2,5 @@ pub mod context;
 pub mod docker;
 pub mod extension;
 pub mod file_utils;
+pub mod sqlx_utils;
 pub mod tembo_config;

tembo-cli/src/cli/sqlx_utils.rs

@@ -0,0 +1,48 @@
+use crate::Result;
+use base64::{engine::general_purpose, Engine as _};
+use spinners::Spinner;
+use spinners::Spinners;
+use sqlx::migrate::Migrator;
+use sqlx::Pool;
+use sqlx::Postgres;
+use std::path::Path;
+use temboclient::models::ConnectionInfo;
+
+pub struct SqlxUtils {}
+
+impl SqlxUtils {
+    // run sqlx migrate
+    pub async fn run_migrations(connection_info: ConnectionInfo, decode: bool) -> Result {
+        let mut sp = Spinner::new(Spinners::Line, "Running SQL migration".into());
+
+        let user: String;
+        let pwd: String;
+
+        if decode {
+            user = SqlxUtils::b64_decode(&connection_info.user);
+            pwd = SqlxUtils::b64_decode(&connection_info.password);
+        } else {
+            user = connection_info.user;
+            pwd = connection_info.password;
+        }
+
+        let connection_string = format!(
+            "postgresql://{}:{}@{}:{}",
+            user, pwd, connection_info.host, connection_info.port
+        );
+
+        let pool = Pool::<Postgres>::connect(connection_string.as_str()).await?;
+
+        let m = Migrator::new(Path::new("./migrations")).await?;
+        m.run(&pool).await?;
+
+        sp.stop_with_message("- SQL migration completed".to_string());
+
+        Ok(())
+    }
+
+    fn b64_decode(b64_encoded: &str) -> String {
+        let bytes = general_purpose::STANDARD.decode(b64_encoded).unwrap();
+        String::from_utf8(bytes).unwrap()
+    }
+}

tembo-cli/src/cmd/apply.rs

@@ -1,13 +1,16 @@
 use crate::{
     cli::{
         context::{get_current_context, Environment, Target},
+        sqlx_utils::SqlxUtils,
         tembo_config,
     },
     Result,
 };
+use anyhow::Error;
 use clap::{ArgMatches, Command};
 use controller::stacks::get_stack;
 use controller::stacks::types::StackType as ControllerStackType;
+use spinners::{Spinner, Spinners};
 use std::{
     collections::HashMap,
     fs::{self},
@@ -18,11 +21,11 @@ use std::{
 use temboclient::{
     apis::{
         configuration::Configuration,
-        instance_api::{create_instance, get_all, put_instance},
+        instance_api::{create_instance, get_all, get_instance, put_instance},
     },
     models::{
-        Cpu, CreateInstance, Extension, ExtensionInstallLocation, Memory, PgConfig, StackType,
-        Storage, TrunkInstall, UpdateInstance,
+        ConnectionInfo, Cpu, CreateInstance, Extension, ExtensionInstallLocation, Memory, PgConfig,
+        StackType, State, Storage, TrunkInstall, UpdateInstance,
     },
 };
 use tokio::runtime::Runtime;
@@ -85,9 +88,16 @@ fn execute_docker() -> Result<()> {
         // Allows DB instance to be ready before running migrations
         sleep(Duration::from_secs(3));
 
+        let conn_info = ConnectionInfo {
+            host: "localhost".to_owned(),
+            pooler_host: Some(Some("localhost-pooler".to_string())),
+            port: 5432,
+            user: "postgres".to_owned(),
+            password: "postgres".to_owned(),
+        };
         Runtime::new()
             .unwrap()
-            .block_on(Docker::run_sqlx_migrate())?;
+            .block_on(SqlxUtils::run_migrations(conn_info, false))?;
     }
 
     // If all of the above was successful, we can print the url to user
@@ -107,11 +117,27 @@ pub fn execute_tembo_cloud(env: Environment) -> Result<()> {
     };
 
     for (_key, value) in instance_settings.iter() {
-        let instance_id = get_instance_id(value.instance_name.clone(), &config, env.clone())?;
-        if let Some(env_instance_id) = instance_id {
+        let mut instance_id = get_instance_id(value.instance_name.clone(), &config, &env)?;
+
+        if let Some(env_instance_id) = instance_id.clone() {
             update_existing_instance(env_instance_id, value, &config, env.clone());
         } else {
-            create_new_instance(value, &config, env.clone());
+            instance_id = create_new_instance(value, &config, env.clone());
+        }
+
+        loop {
+            let mut sp = Spinner::new(Spinners::Line, "Waiting for instance to be up!".into());
+            sleep(Duration::from_secs(10));
+
+            let connection_info: Option<Box<ConnectionInfo>> =
+                is_instance_up(instance_id.as_ref().unwrap().clone(), &config, &env)?;
+            if connection_info.is_some() {
+                Runtime::new()
+                    .unwrap()
+                    .block_on(SqlxUtils::run_migrations(*connection_info.unwrap(), true))?;
+                sp.stop_with_message("- Instance is now up!".to_string());
+                break;
+            }
         }
     }
 
@@ -121,7 +147,7 @@ pub fn execute_tembo_cloud(env: Environment) -> Result<()> {
 pub fn get_instance_id(
     instance_name: String,
     config: &Configuration,
-    env: Environment,
+    env: &Environment,
 ) -> Result<Option<String>> {
     let v = Runtime::new()
         .unwrap()
@@ -142,6 +168,32 @@ pub fn get_instance_id(
     Ok(None)
 }
 
+pub fn is_instance_up(
+    instance_id: String,
+    config: &Configuration,
+    env: &Environment,
+) -> Result<Option<Box<ConnectionInfo>>> {
+    let v = Runtime::new().unwrap().block_on(get_instance(
+        config,
+        env.org_id.clone().unwrap().as_str(),
+        &instance_id,
+    ));
+
+    match v {
+        Ok(result) => {
+            if result.state == State::Up {
+                return Ok(result.connection_info.unwrap());
+            }
+        }
+        Err(error) => {
+            eprintln!("Error getting instance: {}", error);
+            return Err(Error::new(error));
+        }
+    };
+
+    Ok(None)
+}
+
 fn update_existing_instance(
     instance_id: String,
     value: &InstanceSettings,
@@ -162,13 +214,17 @@ fn update_existing_instance(
             println!(
                 "Instance update started for Instance Id: {}",
                 result.instance_id
-            )
+            );
         }
         Err(error) => eprintln!("Error updating instance: {}", error),
     };
 }
 
-fn create_new_instance(value: &InstanceSettings, config: &Configuration, env: Environment) {
+fn create_new_instance(
+    value: &InstanceSettings,
+    config: &Configuration,
+    env: Environment,
+) -> Option<String> {
     let instance = get_create_instance(value);
 
     let v = Runtime::new().unwrap().block_on(create_instance(
@@ -183,9 +239,15 @@ fn create_new_instance(value: &InstanceSettings, config: &Configuration, env: En
                 "Instance creation started for instance_name: {} with instance_id: {}",
                 result.instance_name, result.instance_id
             );
+
+            return Some(result.instance_id);
+        }
+        Err(error) => {
+            eprintln!("Error creating instance: {}", error);
         }
-        Err(error) => eprintln!("Error creating instance: {}", error),
     };
+
+    None
 }
 
 fn get_create_instance(instance_settings: &InstanceSettings) -> CreateInstance {
@@ -380,6 +442,7 @@ pub fn get_rendered_migrations_file(
     let mut tera = Tera::new("templates/**/*").unwrap();
     let _ = tera.add_raw_template("migrations", &contents);
     let mut context = tera::Context::new();
+
     for (_key, value) in instance_settings.iter() {
         let stack_type = ControllerStackType::from_str(value.stack_type.as_str())
             .unwrap_or(ControllerStackType::Standard);

tembo-cli/src/cmd/delete.rs

@@ -47,7 +47,7 @@ fn execute_tembo_cloud(env: Environment) -> Result<()> {
     };
 
     for (_key, value) in instance_settings.iter() {
-        let instance_id = get_instance_id(value.instance_name.clone(), &config, env.clone())?;
+        let instance_id = get_instance_id(value.instance_name.clone(), &config, &env)?;
         if let Some(env_instance_id) = instance_id {
             let v = Runtime::new().unwrap().block_on(delete_instance(
                 &config,