Merge pull request #150 from tower/develop

v0.3.39 release

Author: bradhe

Cargo.lock

@@ -4,7 +4,7 @@ resolver = "2"
 
 [workspace.package]
 edition = "2021"
-version = "0.3.38"
+version = "0.3.39"
 description = "Tower is the best way to host Python data apps in production"
 rust-version = "1.81"
 authors = ["Brad Heller <brad@tower.dev>"]

Cargo.toml

@@ -7,9 +7,10 @@ rust-version = { workspace = true }
 license = { workspace = true }
 
 [dependencies]
+base64 = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true }
-dirs = { workspace = true } 
+dirs = { workspace = true }
 futures = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }

crates/config/Cargo.toml

@@ -181,17 +181,15 @@ impl Config {
 
         configuration.base_path = base_path.to_string();
 
-        if let Some(session) = &self.session {
+        // Always read from disk to pick up team switches
+        if let Ok(session) = Session::from_config_dir() {
             if let Some(active_team) = &session.active_team {
-                // Use the active team's JWT token
                 configuration.bearer_access_token = Some(active_team.token.jwt.clone());
             } else {
-                // Fall back to session token if no active team
                 configuration.bearer_access_token = Some(session.token.jwt.clone());
             }
         }
 
-        // Store the configuration in self
         configuration
     }
 }

crates/config/src/lib.rs

@@ -1,3 +1,4 @@
+use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
 use chrono::{DateTime, TimeZone, Utc};
 use serde::{Deserialize, Serialize};
 use std::fs;
@@ -9,6 +10,20 @@ use crate::error::Error;
 use tower_api::apis::default_api::describe_session;
 use tower_telemetry::debug;
 
+/// Extracts the account ID (aid) from a Tower JWT token.
+/// Returns None if the JWT is malformed or doesn't contain an aid.
+fn extract_aid_from_jwt(jwt: &str) -> Option<String> {
+    let parts: Vec<&str> = jwt.split('.').collect();
+    if parts.len() != 3 {
+        return None;
+    }
+
+    let payload = parts[1];
+    let decoded = URL_SAFE_NO_PAD.decode(payload).ok()?;
+    let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?;
+    json.get("https://tower.dev/aid")?.as_str().map(String::from)
+}
+
 const DEFAULT_TOWER_URL: &str = "https://api.tower.dev";
 
 pub fn default_tower_url() -> Url {
@@ -163,6 +178,22 @@ impl Session {
         }
     }
 
+    /// Sets the active team based on an account ID (aid) extracted from a JWT.
+    /// Returns true if a matching team was found and set as active, false otherwise.
+    pub fn set_active_team_by_aid(&mut self, aid: &str) -> bool {
+        // Find the team whose JWT contains the matching aid
+        if let Some(team) = self
+            .teams
+            .iter()
+            .find(|team| extract_aid_from_jwt(&team.token.jwt).as_deref() == Some(aid))
+        {
+            self.active_team = Some(team.clone());
+            true
+        } else {
+            false
+        }
+    }
+
     /// Updates the session with data from the API response
     pub fn update_from_api_response(
         &mut self,
@@ -263,36 +294,30 @@ impl Session {
     }
 
     pub fn from_jwt(jwt: &str) -> Result<Self, Error> {
-        // We need to instantiate our own configuration object here, instead of the typical thing
-        // that we do which is turn a Config into a Configuration.
+        let jwt_aid = extract_aid_from_jwt(jwt);
+
         let mut config = tower_api::apis::configuration::Configuration::new();
         config.bearer_access_token = Some(jwt.to_string());
 
-        // We only pull TOWER_URL out of the environment here because we only ever use the JWT and
-        // all that in programmatic contexts (when TOWER_URL is set).
-        let tower_url = if let Ok(val) = std::env::var("TOWER_URL") {
-            val
-        } else {
-            DEFAULT_TOWER_URL.to_string()
-        };
-
-        // Setup the base path to point to the /v1 API endpoint as expected.
+        let tower_url = std::env::var("TOWER_URL").unwrap_or(DEFAULT_TOWER_URL.to_string());
         let mut base_path = Url::parse(&tower_url).unwrap();
         base_path.set_path("/v1");
-
         config.base_path = base_path.to_string();
 
-        // This is a bit of a hairy thing: I didn't want to pull in too much from the Tower API
-        // client, so we're using the raw bindings here.
         match run_future_sync(describe_session(&config)) {
             Ok(resp) => {
-                // Now we need to extract the session from the response.
                 let entity = resp.entity.unwrap();
 
                 match entity {
                     tower_api::apis::default_api::DescribeSessionSuccess::Status200(resp) => {
                         let mut session = Session::from_api_session(&resp.session);
                         session.tower_url = base_path;
+
+                        if let Some(aid) = jwt_aid {
+                            session.set_active_team_by_aid(&aid);
+                        }
+
+                        session.save()?;
                         Ok(session)
                     }
                     tower_api::apis::default_api::DescribeSessionSuccess::UnknownValue(val) => {

crates/config/src/session.rs

@@ -133,36 +133,27 @@ pub async fn do_show(config: Config, cmd: &ArgMatches) {
 
             output::table(headers, rows, Some(&app_response));
         }
-        Err(err) => {
-            output::tower_error(err);
-        }
+        Err(err) => output::tower_error_and_die(err, "Fetching app details failed"),
     }
 }
 
 pub async fn do_list_apps(config: Config) {
-    let resp = api::list_apps(&config).await;
-
-    match resp {
-        Ok(resp) => {
-            let items = resp
-                .apps
-                .iter()
-                .map(|app_summary| {
-                    let app = &app_summary.app;
-                    let desc = if app.short_description.is_empty() {
-                        output::placeholder("No description")
-                    } else {
-                        app.short_description.to_string()
-                    };
-                    format!("{}\n{}", output::title(&app.name), desc)
-                })
-                .collect();
-            output::list(items, Some(&resp.apps));
-        }
-        Err(err) => {
-            output::tower_error(err);
-        }
-    }
+    let resp = output::with_spinner("Listing apps", api::list_apps(&config)).await;
+
+    let items = resp
+        .apps
+        .iter()
+        .map(|app_summary| {
+            let app = &app_summary.app;
+            let desc = if app.short_description.is_empty() {
+                output::placeholder("No description")
+            } else {
+                app.short_description.to_string()
+            };
+            format!("{}\n{}", output::title(&app.name), desc)
+        })
+        .collect();
+    output::list(items, Some(&resp.apps));
 }
 
 pub async fn do_create(config: Config, args: &ArgMatches) {
@@ -172,30 +163,16 @@ pub async fn do_create(config: Config, args: &ArgMatches) {
 
     let description = args.get_one::<String>("description").unwrap();
 
-    let mut spinner = output::spinner("Creating app");
+    let app =
+        output::with_spinner("Creating app", api::create_app(&config, name, description)).await;
 
-    match api::create_app(&config, name, description).await {
-        Ok(app) => {
-            spinner.success();
-            output::success_with_data(&format!("App '{}' created", name), Some(app));
-        }
-        Err(err) => {
-            spinner.failure();
-            output::tower_error(err);
-        }
-    }
+    output::success_with_data(&format!("App '{}' created", name), Some(app));
 }
 
 pub async fn do_delete(config: Config, cmd: &ArgMatches) {
     let name = extract_app_name("delete", cmd.subcommand());
-    let mut spinner = output::spinner("Deleting app");
 
-    if let Err(err) = api::delete_app(&config, &name).await {
-        spinner.failure();
-        output::tower_error(err);
-    } else {
-        spinner.success();
-    }
+    output::with_spinner("Deleting app", api::delete_app(&config, &name)).await;
 }
 
 /// Extract app name and run number from command

crates/tower-cmd/src/apps.rs

@@ -40,11 +40,21 @@ pub async fn do_deploy(config: Config, args: &ArgMatches) {
     let create_app = args.get_flag("create");
     if let Err(err) = deploy_from_dir(config, dir, create_app).await {
         match err {
-            crate::Error::ApiDeployError { source } => output::tower_error(source),
-            crate::Error::ApiDescribeAppError { source } => output::tower_error(source),
-            crate::Error::PackageError { source } => output::package_error(source),
-            crate::Error::TowerfileLoadFailed { source, .. } => output::config_error(source),
-            _ => output::error(&err.to_string()),
+            crate::Error::ApiDeployError { source } => {
+                output::tower_error_and_die(source, "Deploying app failed")
+            }
+            crate::Error::ApiDescribeAppError { source } => {
+                output::tower_error_and_die(source, "Fetching app details failed")
+            }
+            crate::Error::PackageError { source } => {
+                output::package_error(source);
+                std::process::exit(1);
+            }
+            crate::Error::TowerfileLoadFailed { source, .. } => {
+                output::config_error(source);
+                std::process::exit(1);
+            }
+            _ => output::die(&err.to_string()),
         }
     }
 }
@@ -62,13 +72,16 @@ pub async fn deploy_from_dir(
     let api_config = config.into();
 
     // Add app existence check before proceeding
-    util::apps::ensure_app_exists(
+    if let Err(err) = util::apps::ensure_app_exists(
         &api_config,
         &towerfile.app.name,
         &towerfile.app.description,
         create_app,
     )
-    .await?;
+    .await
+    {
+        return Err(crate::Error::ApiDescribeAppError { source: err });
+    }
 
     let spec = PackageSpec::from_towerfile(&towerfile);
     let mut spinner = output::spinner("Building package...");