From 3c9ed442bf4387f3600b6d8847ebac40eef6f68a Mon Sep 17 00:00:00 2001
From: Eddie A Tejeda <669988+eddietejeda@users.noreply.github.com>
Date: Sat, 23 May 2026 09:55:39 -0700
Subject: [PATCH] feat(query): add --database flag and improve post-load UX

Add --database / -d flag to `hotdata query` so users can scope a query
to a specific managed database without permanently changing config via
`databases set`. Adds ApiClient::with_database() builder to override the
X-Database-Id header for a single request.

After `databases create` and `databases load` succeed, print ready-to-run
examples showing the correct `--database <id>` flag, `default.<schema>.<table>`
SQL prefix, and a reminder that uppercase column names need double-quoting.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/api.rs       |  6 ++++++
 src/command.rs   |  4 ++++
 src/databases.rs | 32 +++++++++++++++++++++++++++++++-
 src/main.rs      | 11 +++++++++--
 src/query.rs     | 13 +++++++++++--
 5 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/src/api.rs b/src/api.rs
index 59e9e09..df3ef71 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -122,6 +122,12 @@ impl ApiClient {
         }
     }
 
+    /// Override the database ID for a single query without touching config.
+    pub fn with_database(mut self, database_id: &str) -> Self {
+        self.database_id = Some(database_id.to_string());
+        self
+    }
+
     /// Test-only client (no config load). Used with a local mock HTTP server.
     #[cfg(test)]
     pub(crate) fn test_new(api_url: &str, api_key: &str, workspace_id: Option<&str>) -> Self {
diff --git a/src/command.rs b/src/command.rs
index 0c1a281..bfc2171 100644
--- a/src/command.rs
+++ b/src/command.rs
@@ -38,6 +38,10 @@ pub enum Commands {
         #[arg(long)]
         connection: Option<String>,
 
+        /// Run query against a specific managed database (overrides the current database set via `databases set`)
+        #[arg(long, short = 'd')]
+        database: Option<String>,
+
         /// Output format
         #[arg(long = "output", short = 'o', default_value = "table", value_parser = ["table", "json", "csv"])]
         output: String,
diff --git a/src/databases.rs b/src/databases.rs
index 542cb08..15526dd 100644
--- a/src/databases.rs
+++ b/src/databases.rs
@@ -463,6 +463,22 @@ pub fn create(
                 println!("description: {desc}");
             }
             println!("id:          {}", result.id);
+            println!();
+            println!(
+                "{}",
+                format!(
+                    concat!(
+                        "Load a table:\n",
+                        "  hotdata databases load --file <path.parquet> {}.<table_name>\n",
+                        "\nQuery with:\n",
+                        "  hotdata query --database {} \"SELECT * FROM default.public.<table> LIMIT 10\"\n",
+                        "\n  Tip: use 'default.<schema>.<table>' as the SQL prefix (not the database or connection id)\n",
+                        "  Column names are case-sensitive — wrap uppercase names in double quotes",
+                    ),
+                    result.id, result.id
+                )
+                .dark_grey()
+            );
         }
         _ => unreachable!(),
     }
@@ -614,8 +630,22 @@ pub fn tables_load(
 
     let full_name = format!("default.{}.{}", result.schema_name, result.table_name);
     println!("{}", "Table loaded".green());
-    println!("full_name: {}", full_name.green());
+    println!("full_name: {}", full_name.clone().green());
     println!("rows:      {}", result.row_count);
+    println!();
+    println!(
+        "{}",
+        format!(
+            concat!(
+                "Query it now:\n",
+                "  hotdata query \"SELECT * FROM {} LIMIT 10\"\n",
+                "\n  Tip: column names are case-sensitive.\n",
+                "  Wrap uppercase names in double quotes: SELECT \"MyColumn\" FROM {} LIMIT 10",
+            ),
+            full_name, full_name
+        )
+        .dark_grey()
+    );
 }
 
 pub fn tables_delete(workspace_id: &str, database: Option<&str>, table: &str, schema: Option<&str>) {
diff --git a/src/main.rs b/src/main.rs
index 66e84ed..7cd9167 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -260,6 +260,7 @@ fn main() {
                 sql,
                 workspace_id,
                 connection,
+                database,
                 output,
                 command,
             } => {
@@ -268,7 +269,13 @@ fn main() {
                     Some(QueryCommands::Status { id }) => query::poll(&id, &workspace_id, &output),
                     None => match sql {
                         Some(sql) => {
-                            query::execute(&sql, &workspace_id, connection.as_deref(), &output)
+                            query::execute(
+                                &sql,
+                                &workspace_id,
+                                connection.as_deref(),
+                                database.as_deref(),
+                                &output,
+                            )
                         }
                         None => {
                             use clap::CommandFactory;
@@ -830,7 +837,7 @@ fn main() {
                     ),
                     _ => unreachable!(),
                 };
-                query::execute(&sql, &workspace_id, None, &output)
+                query::execute(&sql, &workspace_id, None, None, &output)
             }
             Commands::Queries {
                 id,
diff --git a/src/query.rs b/src/query.rs
index 2e70ad7..82d275c 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -45,8 +45,17 @@ fn value_to_string(v: &Value) -> String {
     }
 }
 
-pub fn execute(sql: &str, workspace_id: &str, connection: Option<&str>, format: &str) {
-    let api = ApiClient::new(Some(workspace_id));
+pub fn execute(
+    sql: &str,
+    workspace_id: &str,
+    connection: Option<&str>,
+    database: Option<&str>,
+    format: &str,
+) {
+    let mut api = ApiClient::new(Some(workspace_id));
+    if let Some(db_id) = database {
+        api = api.with_database(db_id);
+    }
 
     let mut body = serde_json::json!({
         "sql": sql,