Merge pull request #254 from hotdata-dev/feat/databases-count

feat(databases): count databases in the workspace

Author: shefeek-jinnah

Cargo.lock

@@ -16,7 +16,7 @@ path = "src/main.rs"
 # behind a shared multi-thread tokio runtime. The `arrow` feature backs result
 # decode; `upload_file` runs the presigned direct-to-storage upload flow
 # (see src/sdk.rs::Api::upload).
-hotdata = { version = "0.10.0", features = ["arrow"] }
+hotdata = { version = "0.11.0", features = ["arrow"] }
 # Shared multi-thread runtime for the sync wrapper; block_on is called
 # concurrently from rayon worker threads (see src/indexes.rs). The presigned
 # upload (src/sdk.rs::Api::upload) drives the SDK's async `upload_file` on this

Cargo.toml

@@ -95,6 +95,7 @@ Returns workspaces with `public_id`, `name`, `active`, `favorite`, `provision_st
 
 ```
 hotdata databases list [--workspace-id <workspace_id>] [--output table|json|yaml]
+hotdata databases count [--workspace-id <workspace_id>] [--output table|json|yaml]
 hotdata databases create [--name <display_name>] [--catalog <alias>] [--table <table> ...] [--schema public] [--expires-at <duration|timestamp>] [--workspace-id <workspace_id>] [--output table|json|yaml]
 hotdata databases fork [<id>] [--name <display_name>] [--expires-at <duration|timestamp>] [--workspace-id <workspace_id>] [--output table|json|yaml]
 hotdata databases set <id>
@@ -118,6 +119,7 @@ hotdata databases tables delete <table> [--database <id>] [--schema public] [--w
 ```
 
 - `list` — all managed databases in the workspace. Active database is marked with `*` under the DEFAULT column; CREATED shows when each database was made.
+- `count` — the total number of managed databases in the workspace, across **all** pages (`list` shows one page). Prints a bare integer by default so it drops straight into scripts (`$(hotdata databases count)`); `--output json|yaml` render `{"count": N}` / `count: N`.
 - `create` — creates a new managed database. `--name` is an optional human-readable display name. `--catalog` sets the SQL alias used in queries (`SELECT … FROM <catalog>.schema.table`); must be `[a-z_][a-z0-9_]*`. `--expires-at` accepts relative durations (`24h`, `7d`, `90m`) or an RFC 3339 timestamp; omitting means no expiry. Repeat `--table` to declare tables up front.
 - `fork` — creates a new managed database that is an independent deep copy of an existing one (same schemas, tables, and data); the source is left unchanged and the two diverge freely afterwards. The source defaults to the active database; pass the database `<id>` to fork another. `--name` defaults to `<source>-fork` (so the two stay distinguishable in `list`); `--expires-at` accepts a relative duration or RFC 3339 timestamp, and when omitted a still-future source expiry is carried over. The fork becomes the active database on success. The fork answers to the **same catalog alias** as its source inside its own scope; catalogs attached to the source are **re-attached** to the fork, but indexes are **not** carried over. Only databases created with the current (DuckLake) storage engine can be forked — older parquet-backed databases return an error.
 - `set` — saves the database **id** as the active database. Subsequent `databases tables` and `context` commands use it automatically. Note that a successful `fork` also updates this: the fork becomes the active database.

skills/hotdata/SKILL.md

@@ -21,6 +21,13 @@ pub enum DatabasesCommands {
         output: String,
     },
 
+    /// Count managed databases in the workspace
+    Count {
+        /// Output format
+        #[arg(long = "output", short = 'o', default_value = "table", value_parser = ["table", "json", "yaml"])]
+        output: String,
+    },
+
     /// Show details for a specific managed database
     Show {
         /// Database name or ID
@@ -440,9 +447,15 @@ fn list_all_databases(
             (true, Some(msg)) => block_with_wakeup(
                 api,
                 msg,
-                api.client().databases().list(None, cursor.as_deref()),
+                api.client()
+                    .databases()
+                    .list(None, cursor.as_deref(), None, None),
+            )?,
+            _ => block(
+                api.client()
+                    .databases()
+                    .list(None, cursor.as_deref(), None, None),
             )?,
-            _ => block(api.client().databases().list(None, cursor.as_deref()))?,
         };
         first = false;
         // Move the page rows out first (partial move), then read next_cursor —
@@ -938,7 +951,7 @@ pub fn list(workspace_id: &str, format: &str, limit: Option<u32>, cursor: Option
         "Loading databases...",
         api.client()
             .databases()
-            .list(limit.map(|l| l as i32), cursor),
+            .list(limit.map(|l| l as i32), cursor, None, None),
     )
     .unwrap_or_else(|e| e.exit());
     let has_more = resp.has_more.flatten().unwrap_or(false);
@@ -1003,6 +1016,54 @@ pub fn list(workspace_id: &str, format: &str, limit: Option<u32>, cursor: Option
     }
 }
 
+/// Serialized shape for `databases count -o json|yaml`.
+///
+/// The field stays `count` even though the endpoint answers with `total`: this
+/// is the CLI's own output contract, and `hotdata databases count` returning
+/// `{"count": N}` is what reads naturally at the command line. The API needs
+/// the other name because its list response already spends `count` on a page
+/// size.
+#[derive(Serialize)]
+struct DatabaseCount {
+    count: i64,
+}
+
+/// Total number of managed databases in the workspace.
+///
+/// One request. The list endpoint's `count` field is a page size, not a
+/// workspace total, so this used to drain every page and take the length —
+/// which cost a round trip per 100 databases and threw every row away on
+/// arrival. `GET /v1/databases/count` answers it in a single indexed query.
+fn count_databases(api: &Api) -> Result<i64, ApiError> {
+    block_with_wakeup(
+        api,
+        "Counting databases...",
+        api.client().databases().count(None, None),
+    )
+    .map(|resp| resp.total)
+}
+
+pub fn count(workspace_id: &str, format: &str) {
+    let api = Api::new(Some(workspace_id));
+    let count = count_databases(&api).unwrap_or_else(|e| e.exit());
+
+    match format {
+        "json" => println!(
+            "{}",
+            serde_json::to_string_pretty(&DatabaseCount { count }).unwrap()
+        ),
+        "yaml" => print!(
+            "{}",
+            serde_yaml::to_string(&DatabaseCount { count }).unwrap()
+        ),
+        // Bare integer, nothing else, so $(hotdata databases count) is directly
+        // usable in scripts. Zero is the answer for an empty workspace — no
+        // "No databases found." hint like `list` has.
+        "table" => println!("{count}"),
+        _ => unreachable!(),
+    }
+}
+
 pub fn get(workspace_id: &str, id_or_name: &str, format: &str) {
     let api = Api::new(Some(workspace_id));
     let db = resolve_database(&api, id_or_name);
@@ -2381,7 +2442,7 @@ mod tests {
             .with_status(200)
             .with_header("content-type", "application/json")
             .with_body(
-                r#"{"count":1,"limit":1000,"tables":[{"connection":"default","schema":"public","table":"b","synced":true,"last_sync":null}],"has_more":false,"next_cursor":null}"#,
+                r#"{"count":1,"limit":1000,"tables":[{"connection":"default","schema":"public","table":"b","synced":true,"partition_by":[],"sorted_by":[],"last_sync":null}],"has_more":false,"next_cursor":null}"#,
             )
             .create();
         let page0 = server
@@ -2393,7 +2454,7 @@ mod tests {
             .with_status(200)
             .with_header("content-type", "application/json")
             .with_body(
-                r#"{"count":1,"limit":1000,"tables":[{"connection":"default","schema":"public","table":"a","synced":false,"last_sync":null}],"has_more":true,"next_cursor":"cur2"}"#,
+                r#"{"count":1,"limit":1000,"tables":[{"connection":"default","schema":"public","table":"a","synced":false,"partition_by":[],"sorted_by":[],"last_sync":null}],"has_more":true,"next_cursor":"cur2"}"#,
             )
             .create();
 
@@ -2438,6 +2499,62 @@ mod tests {
         assert_eq!(dbs[1].id, "db2");
     }
 
+    #[test]
+    fn count_databases_reads_the_total_from_the_count_endpoint() {
+        let mut server = mockito::Server::new();
+        // The dedicated endpoint, hit exactly once — no paging. `expect(1)` is
+        // the assertion that matters: the old implementation drained
+        // /v1/databases page by page, and a regression back to that would show
+        // up here as the wrong request being made.
+        let count = server
+            .mock("GET", "/v1/databases/count")
+            .with_status(200)
+            .with_header("content-type", "application/json")
+            .with_body(r#"{"total":7}"#)
+            .expect(1)
+            .create();
+        // Registered but never expected to be called: counting must not touch
+        // the listing at all.
+        let listing = server
+            .mock("GET", "/v1/databases")
+            .with_status(200)
+            .with_header("content-type", "application/json")
+            .with_body(
+                r#"{"databases":[],"count":0,"limit":100,"has_more":false,"next_cursor":null}"#,
+            )
+            .expect(0)
+            .create();
+
+        let api = Api::test_new(&server.url(), "k", Some("ws"));
+        let n = count_databases(&api).unwrap();
+        count.assert();
+        listing.assert();
+        assert_eq!(n, 7);
+    }
+
+    #[test]
+    fn count_databases_empty_workspace_is_zero() {
+        let mut server = mockito::Server::new();
+        let count = server
+            .mock("GET", "/v1/databases/count")
+            .with_status(200)
+            .with_header("content-type", "application/json")
+            .with_body(r#"{"total":0}"#)
+            .create();
+
+        let api = Api::test_new(&server.url(), "k", Some("ws"));
+        let n = count_databases(&api).unwrap();
+        count.assert();
+        assert_eq!(n, 0);
+    }
+
+    #[test]
+    fn database_count_serialization_shapes() {
+        let c = DatabaseCount { count: 7 };
+        assert_eq!(serde_json::to_string(&c).unwrap(), r#"{"count":7}"#);
+        assert_eq!(serde_yaml::to_string(&c).unwrap(), "count: 7\n");
+    }
+
     #[test]
     fn collect_tables_with_limit_makes_single_request_and_returns_cursor() {
         let mut server = mockito::Server::new();
@@ -2451,7 +2568,7 @@ mod tests {
             .with_status(200)
             .with_header("content-type", "application/json")
             .with_body(
-                r#"{"count":1,"limit":5,"tables":[{"connection":"default","schema":"public","table":"a","synced":true,"last_sync":null}],"has_more":true,"next_cursor":"page2"}"#,
+                r#"{"count":1,"limit":5,"tables":[{"connection":"default","schema":"public","table":"a","synced":true,"partition_by":[],"sorted_by":[],"last_sync":null}],"has_more":true,"next_cursor":"page2"}"#,
             )
             .expect(1)
             .create();
@@ -2477,7 +2594,7 @@ mod tests {
             .with_status(200)
             .with_header("content-type", "application/json")
             .with_body(
-                r#"{"count":1,"limit":1000,"tables":[{"connection":"default","schema":"public","table":"orders","synced":true,"last_sync":null}],"has_more":false,"next_cursor":null}"#,
+                r#"{"count":1,"limit":1000,"tables":[{"connection":"default","schema":"public","table":"orders","synced":true,"partition_by":[],"sorted_by":[],"last_sync":null}],"has_more":false,"next_cursor":null}"#,
             )
             .create();
 

src/commands/databases.rs

@@ -855,7 +855,7 @@ mod tests {
             .with_header("content-type", "application/json")
             .with_body(
                 r#"{"count":1,"limit":100,"tables":[
-                {"connection":"__db_abc","schema":"public","table":"vec_mid","synced":true}
+                {"connection":"__db_abc","schema":"public","table":"vec_mid","synced":true,"partition_by":[],"sorted_by":[]}
             ],"has_more":false,"next_cursor":null}"#,
             )
             .create();
@@ -937,7 +937,7 @@ mod tests {
             .with_header("content-type", "application/json")
             .with_body(
                 r#"{"count":1,"limit":100,"tables":[
-                {"connection":"__db_abc","schema":"public","table":"listings","synced":true}
+                {"connection":"__db_abc","schema":"public","table":"listings","synced":true,"partition_by":[],"sorted_by":[]}
             ],"has_more":false,"next_cursor":null}"#,
             )
             .create();
@@ -992,7 +992,7 @@ mod tests {
             .with_header("content-type", "application/json")
             .with_body(
                 r#"{"count":1,"limit":100,"tables":[
-                {"connection":"Warehouse","schema":"public","table":"events","synced":true}
+                {"connection":"Warehouse","schema":"public","table":"events","synced":true,"partition_by":[],"sorted_by":[]}
             ],"has_more":false,"next_cursor":null}"#,
             )
             .create();
@@ -1036,7 +1036,7 @@ mod tests {
             .with_header("content-type", "application/json")
             .with_body(
                 r#"{"count":1,"limit":100,"tables":[
-                {"connection":"__db_abc","schema":"public","table":"listings","synced":true}
+                {"connection":"__db_abc","schema":"public","table":"listings","synced":true,"partition_by":[],"sorted_by":[]}
             ],"has_more":false,"next_cursor":null}"#,
             )
             .create();
@@ -1082,8 +1082,8 @@ mod tests {
             .with_header("content-type", "application/json")
             .with_body(
                 r#"{"count":2,"limit":100,"tables":[
-                {"connection":"c1","schema":"public","table":"z","synced":true},
-                {"connection":"c1","schema":"public","table":"a","synced":true}
+                {"connection":"c1","schema":"public","table":"z","synced":true,"partition_by":[],"sorted_by":[]},
+                {"connection":"c1","schema":"public","table":"a","synced":true,"partition_by":[],"sorted_by":[]}
             ],"has_more":false,"next_cursor":null}"#,
             )
             .create();

src/commands/indexes.rs

@@ -262,6 +262,9 @@ fn main() {
                             limit,
                             cursor,
                         }) => databases::list(&workspace_id, &output, limit, cursor.as_deref()),
+                        Some(DatabasesCommands::Count { output }) => {
+                            databases::count(&workspace_id, &output)
+                        }
                         Some(DatabasesCommands::Show { name_or_id, output }) => {
                             databases::get(&workspace_id, &name_or_id, &output)
                         }

src/main.rs

@@ -14,13 +14,29 @@ fn databases_help_lists_subcommands() {
     );
     let help = String::from_utf8_lossy(&output.stdout);
     assert!(help.contains("list"));
+    assert!(help.contains("count"));
     assert!(help.contains("create"));
     assert!(help.contains("delete"));
     assert!(help.contains("tables"));
     assert!(help.contains("attach"));
     assert!(help.contains("detach"));
 }
 
+#[test]
+fn databases_count_help_documents_output_flag() {
+    let output = hotdata()
+        .args(["databases", "count", "--help"])
+        .output()
+        .unwrap();
+    assert!(
+        output.status.success(),
+        "stderr: {}",
+        String::from_utf8_lossy(&output.stderr)
+    );
+    let help = String::from_utf8_lossy(&output.stdout);
+    assert!(help.contains("--output"), "help: {help}");
+}
+
 #[test]
 fn databases_create_help_documents_attach_flag() {
     let output = hotdata()