Merge main into feat/datasets-rework

eddietejeda · eddietejeda · commit ee62435439d9 · 2026-05-22T13:23:57.000-07:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -40,6 +40,7 @@ sysinfo = { version = "0.38.4", default-features = false, features = ["system"]
 self_update = { version = "0.42", default-features = false, features = ["rustls"] }
 lzma-rs = "0.3"
 tempfile = "3"
+urlencoding = "2.1.3"
 
 [dev-dependencies]
 mockito = "1"
diff --git a/src/api.rs b/src/api.rs
@@ -32,6 +32,7 @@ pub struct ApiClient {
     pub api_url: String,
     workspace_id: Option<String>,
     sandbox_id: Option<String>,
+    database_id: Option<String>,
 }
 
 impl ApiClient {
@@ -117,6 +118,7 @@ impl ApiClient {
                 }
                 profile_config.sandbox
             }),
+            database_id: workspace_id.and_then(|ws| crate::config::load_current_database("default", ws)),
         }
     }
 
@@ -129,6 +131,7 @@ impl ApiClient {
             api_url: api_url.to_string(),
             workspace_id: workspace_id.map(String::from),
             sandbox_id: None,
+            database_id: None,
         }
     }
 
@@ -167,6 +170,9 @@ impl ApiClient {
             req = req.header("X-Session-Id", sid);
             req = req.header("X-Sandbox-Id", sid);
         }
+        if let Some(ref db_id) = self.database_id {
+            req = req.header("X-Database-Id", db_id);
+        }
         req
     }
 
diff --git a/src/command.rs b/src/command.rs
@@ -558,6 +558,12 @@ pub enum DatabasesCommands {
         output: String,
     },
 
+    /// Set the current database (used by default when no database is specified)
+    Set {
+        /// Database id or description
+        id_or_description: String,
+    },
+
     /// Delete a managed database and its tables
     Delete {
         /// Database name or connection ID
@@ -594,8 +600,9 @@ pub enum DatabasesCommands {
 pub enum DatabaseTablesCommands {
     /// List tables in a managed database
     List {
-        /// Database name or connection ID
-        database: String,
+        /// Database id or description (defaults to current database)
+        #[arg(long)]
+        database: Option<String>,
 
         /// Filter by schema name
         #[arg(long)]
@@ -608,8 +615,9 @@ pub enum DatabaseTablesCommands {
 
     /// Load a parquet file into a table (creates or replaces the table)
     Load {
-        /// Database name or connection ID
-        database: String,
+        /// Database id or description (defaults to current database)
+        #[arg(long)]
+        database: Option<String>,
 
         /// Table name
         table: String,
@@ -633,8 +641,9 @@ pub enum DatabaseTablesCommands {
 
     /// Delete a table from a managed database
     Delete {
-        /// Database name or connection ID
-        database: String,
+        /// Database id or description (defaults to current database)
+        #[arg(long)]
+        database: Option<String>,
 
         /// Table name
         table: String,
diff --git a/src/config.rs b/src/config.rs
@@ -101,6 +101,8 @@ pub struct ProfileConfig {
     pub workspaces: Vec<WorkspaceEntry>,
     #[serde(default, skip_serializing_if = "Option::is_none", alias = "session")]
     pub sandbox: Option<String>,
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub current_databases: HashMap<String, String>,
 }
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -227,6 +229,60 @@ pub fn clear_sandbox(profile: &str) -> Result<(), String> {
     write_config(&config_path, &content)
 }
 
+pub fn save_current_database(profile: &str, workspace_id: &str, database_id: &str) -> Result<(), String> {
+    let config_path = config_path()?;
+
+    let mut config_file: ConfigFile = if config_path.exists() {
+        let content = fs::read_to_string(&config_path)
+            .map_err(|e| format!("error reading config file: {e}"))?;
+        serde_yaml::from_str(&content).map_err(|e| format!("error parsing config file: {e}"))?
+    } else {
+        ConfigFile { profiles: HashMap::new() }
+    };
+
+    config_file
+        .profiles
+        .entry(profile.to_string())
+        .or_default()
+        .current_databases
+        .insert(workspace_id.to_string(), database_id.to_string());
+
+    let content = serde_yaml::to_string(&config_file)
+        .map_err(|e| format!("error serializing config: {e}"))?;
+    write_config(&config_path, &content)
+}
+
+pub fn load_current_database(profile: &str, workspace_id: &str) -> Option<String> {
+    let config_path = config_path().ok()?;
+    if !config_path.exists() {
+        return None;
+    }
+    let content = fs::read_to_string(&config_path).ok()?;
+    let config_file: ConfigFile = serde_yaml::from_str(&content).ok()?;
+    config_file.profiles.get(profile)?.current_databases.get(workspace_id).cloned()
+}
+
+pub fn clear_current_database(profile: &str, workspace_id: &str) -> Result<(), String> {
+    let config_path = config_path()?;
+
+    if !config_path.exists() {
+        return Ok(());
+    }
+
+    let content = fs::read_to_string(&config_path)
+        .map_err(|e| format!("error reading config file: {e}"))?;
+    let mut config_file: ConfigFile =
+        serde_yaml::from_str(&content).map_err(|e| format!("error parsing config file: {e}"))?;
+
+    if let Some(entry) = config_file.profiles.get_mut(profile) {
+        entry.current_databases.remove(workspace_id);
+    }
+
+    let content = serde_yaml::to_string(&config_file)
+        .map_err(|e| format!("error serializing config: {e}"))?;
+    write_config(&config_path, &content)
+}
+
 pub fn resolve_workspace_id(provided: Option<String>, profile_config: &ProfileConfig) -> Result<String, String> {
     if let Some(id) = provided {
         return Ok(id);
diff --git a/src/databases.rs b/src/databases.rs
@@ -23,6 +23,7 @@ pub struct Database {
     pub id: String,
     pub description: Option<String>,
     pub default_connection_id: String,
+    #[serde(default)]
     attachments: Vec<DatabaseAttachment>,
 }
 
@@ -82,7 +83,10 @@ fn fetch_database(api: &ApiClient, id: &str) -> Database {
 
 pub fn try_resolve_database(api: &ApiClient, id_or_description: &str) -> Result<Database, String> {
     // Try a direct id lookup first — avoids the list round-trip for the common case.
-    if let Some(db) = api.get_none_if_not_found(&format!("/databases/{id_or_description}")) {
+    // Percent-encode the segment so descriptions containing spaces or other URL-unsafe
+    // characters don't cause a URL parse error before the list fallback can run.
+    let encoded = urlencoding::encode(id_or_description);
+    if let Some(db) = api.get_none_if_not_found(&format!("/databases/{encoded}")) {
         return Ok(db);
     }
 
@@ -436,6 +440,11 @@ pub fn create(
         }
     };
 
+    if let Err(e) = crate::config::save_current_database("default", workspace_id, &result.id) {
+        use crossterm::style::Stylize;
+        eprintln!("{}", format!("warning: database created but could not set as current: {e}").yellow());
+    }
+
     match format {
         "json" => println!("{}", serde_json::to_string_pretty(&result).unwrap()),
         "yaml" => print!("{}", serde_yaml::to_string(&result).unwrap()),
@@ -450,6 +459,34 @@ pub fn create(
     }
 }
 
+pub fn set(workspace_id: &str, id_or_description: &str) {
+    use crossterm::style::Stylize;
+    let api = ApiClient::new(Some(workspace_id));
+    let db = resolve_database(&api, id_or_description);
+    if let Err(e) = crate::config::save_current_database("default", workspace_id, &db.id) {
+        eprintln!("{}", format!("error saving current database: {e}").red());
+        std::process::exit(1);
+    }
+    println!("{}", format!("Current database set to {}", db.id).green());
+}
+
+fn resolve_current_database(provided: Option<&str>, workspace_id: &str) -> String {
+    if let Some(id) = provided {
+        return id.to_string();
+    }
+    match crate::config::load_current_database("default", workspace_id) {
+        Some(id) => id,
+        None => {
+            use crossterm::style::Stylize;
+            eprintln!(
+                "{}",
+                "error: no current database set. Use 'hotdata databases set <id>' or pass a database id.".red()
+            );
+            std::process::exit(1);
+        }
+    }
+}
+
 pub fn delete(workspace_id: &str, id_or_description: &str) {
     use crossterm::style::Stylize;
 
@@ -462,12 +499,19 @@ pub fn delete(workspace_id: &str, id_or_description: &str) {
         std::process::exit(1);
     }
 
+    // If the deleted database was the current one, clear it so subsequent
+    // commands don't silently send a stale X-Database-Id header.
+    if crate::config::load_current_database("default", workspace_id).as_deref() == Some(&db.id) {
+        let _ = crate::config::clear_current_database("default", workspace_id);
+    }
+
     println!("{}", "Database deleted.".green());
 }
 
-pub fn tables_list(workspace_id: &str, database: &str, schema: Option<&str>, format: &str) {
+pub fn tables_list(workspace_id: &str, database: Option<&str>, schema: Option<&str>, format: &str) {
+    let database = resolve_current_database(database, workspace_id);
     let api = ApiClient::new(Some(workspace_id));
-    let db = resolve_database(&api, database);
+    let db = resolve_database(&api, &database);
     let tables = collect_tables(&api, &db.default_connection_id, schema);
 
     let rows = table_rows(tables);
@@ -502,7 +546,7 @@ pub fn tables_list(workspace_id: &str, database: &str, schema: Option<&str>, for
 
 pub fn tables_load(
     workspace_id: &str,
-    database: &str,
+    database: Option<&str>,
     table: &str,
     schema: Option<&str>,
     file: Option<&str>,
@@ -511,8 +555,9 @@ pub fn tables_load(
 ) {
     use crossterm::style::Stylize;
 
+    let database = resolve_current_database(database, workspace_id);
     let api = ApiClient::new(Some(workspace_id));
-    let db = resolve_database(&api, database);
+    let db = resolve_database(&api, &database);
     let schema = schema_name(schema);
 
     // clap enforces mutual exclusion; only one of these is ever Some.
@@ -564,11 +609,12 @@ pub fn tables_load(
     println!("rows:      {}", result.row_count);
 }
 
-pub fn tables_delete(workspace_id: &str, database: &str, table: &str, schema: Option<&str>) {
+pub fn tables_delete(workspace_id: &str, database: Option<&str>, table: &str, schema: Option<&str>) {
     use crossterm::style::Stylize;
 
+    let database = resolve_current_database(database, workspace_id);
     let api = ApiClient::new(Some(workspace_id));
-    let db = resolve_database(&api, database);
+    let db = resolve_database(&api, &database);
     let schema = schema_name(schema);
 
     let path = managed_table_delete_path(&db.default_connection_id, schema, table);
diff --git a/src/main.rs b/src/main.rs
@@ -56,6 +56,10 @@ struct Cli {
     command: Option<Commands>,
 }
 
+/// Set once after workspace resolution so the database footer can reference it
+/// without re-doing config I/O.
+static ACTIVE_WORKSPACE_ID: std::sync::OnceLock<String> = std::sync::OnceLock::new();
+
 fn resolve_workspace(provided: Option<String>) -> String {
     // HOTDATA_WORKSPACE env var takes priority and blocks --workspace-id flag
     if let Ok(ws) = std::env::var("HOTDATA_WORKSPACE") {
@@ -67,6 +71,7 @@ fn resolve_workspace(provided: Option<String>) -> String {
             );
             std::process::exit(1);
         }
+        let _ = ACTIVE_WORKSPACE_ID.set(ws.clone());
         return ws;
     }
     if sandbox::find_sandbox_run_ancestor().is_some() {
@@ -75,7 +80,10 @@ fn resolve_workspace(provided: Option<String>) -> String {
     }
     match config::load("default") {
         Ok(profile) => match config::resolve_workspace_id(provided, &profile) {
-            Ok(id) => id,
+            Ok(id) => {
+                let _ = ACTIVE_WORKSPACE_ID.set(id.clone());
+                id
+            }
             Err(e) => {
                 eprintln!("error: {e}");
                 std::process::exit(1);
@@ -125,11 +133,25 @@ extern "C" fn print_sandbox_footer() {
     );
 }
 
+extern "C" fn print_database_footer() {
+    use crossterm::style::Stylize;
+    if let Some(ws_id) = ACTIVE_WORKSPACE_ID.get() {
+        if let Some(id) = config::load_current_database("default", ws_id) {
+            eprintln!(
+                "{}",
+                format!("current database: {id}  use 'hotdata databases set' to change")
+                    .dark_grey(),
+            );
+        }
+    }
+}
+
 fn main() {
     // Register before `Cli::parse`, since `--help` / `--version` exit
     // from inside the parser. Safety: `atexit` is async-signal-safe;
     // the callback only reads env vars / files and writes to stderr.
     unsafe { atexit(print_sandbox_footer) };
+    unsafe { atexit(print_database_footer) };
 
     dotenvy::dotenv().ok();
     let cli = Cli::parse();
@@ -373,6 +395,9 @@ fn main() {
                             &tables,
                             &output,
                         ),
+                        Some(DatabasesCommands::Set { id_or_description }) => {
+                            databases::set(&workspace_id, &id_or_description)
+                        }
                         Some(DatabasesCommands::Delete { name_or_id }) => {
                             databases::delete(&workspace_id, &name_or_id)
                         }
@@ -385,7 +410,7 @@ fn main() {
                             let (database, schema, table) = parse_db_target(&target);
                             databases::tables_load(
                                 &workspace_id,
-                                &database,
+                                Some(database.as_str()),
                                 &table,
                                 Some(schema.as_str()),
                                 file.as_deref(),
@@ -400,7 +425,7 @@ fn main() {
                                 output,
                             } => databases::tables_list(
                                 &workspace_id,
-                                &database,
+                                database.as_deref(),
                                 schema.as_deref(),
                                 &output,
                             ),
@@ -413,7 +438,7 @@ fn main() {
                                 upload_id,
                             } => databases::tables_load(
                                 &workspace_id,
-                                &database,
+                                database.as_deref(),
                                 &table,
                                 Some(schema.as_str()),
                                 file.as_deref(),
@@ -426,7 +451,7 @@ fn main() {
                                 schema,
                             } => databases::tables_delete(
                                 &workspace_id,
-                                &database,
+                                database.as_deref(),
                                 &table,
                                 Some(schema.as_str()),
                             ),
diff --git a/tests/databases_cli.rs b/tests/databases_cli.rs

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@ pub struct ApiClient {`
`32`	`32`	`pub api_url: String,`
`33`	`33`	`workspace_id: Option<String>,`
`34`	`34`	`sandbox_id: Option<String>,`
	`35`	`+ database_id: Option<String>,`
`35`	`36`	`}`
`36`	`37`
`37`	`38`	`impl ApiClient {`
`@@ -117,6 +118,7 @@ impl ApiClient {`
`117`	`118`	`}`
`118`	`119`	`profile_config.sandbox`
`119`	`120`	`}),`
	`121`	`+ database_id: workspace_id.and_then(\|ws\| crate::config::load_current_database("default", ws)),`
`120`	`122`	`}`
`121`	`123`	`}`
`122`	`124`
`@@ -129,6 +131,7 @@ impl ApiClient {`
`129`	`131`	`api_url: api_url.to_string(),`
`130`	`132`	`workspace_id: workspace_id.map(String::from),`
`131`	`133`	`sandbox_id: None,`
	`134`	`+ database_id: None,`
`132`	`135`	`}`
`133`	`136`	`}`
`134`	`137`
`@@ -167,6 +170,9 @@ impl ApiClient {`
`167`	`170`	`req = req.header("X-Session-Id", sid);`
`168`	`171`	`req = req.header("X-Sandbox-Id", sid);`
`169`	`172`	`}`
	`173`	`+ if let Some(ref db_id) = self.database_id {`
	`174`	`+ req = req.header("X-Database-Id", db_id);`
	`175`	`+ }`
`170`	`176`	`req`
`171`	`177`	`}`
`172`	`178`