From a21fab1f0ccaa652e93d5caad3a8f39407044161 Mon Sep 17 00:00:00 2001 From: Shefeek Jinnah Date: Wed, 12 Aug 2026 18:43:08 +0530 Subject: [PATCH] feat(databases): count databases in the workspace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `hotdata databases count` reports how many managed databases a workspace holds. Prints a bare integer by default so it composes in scripts; --output json/yaml render {"count": N} / count: N. `databases list` shows one page on purpose, and the API's per-page `count` field is a page size, not a total, so there was no way to ask the question without paging by hand and adding up rows. Backed by GET /v1/databases/count (runtimedb #1053, SDK 0.11.0), which answers from a single indexed aggregate. The alternative — draining the paged listing client-side — costs one sequential round trip per 100 databases and discards every row on arrival. Also in this change, both forced by the 0.11.0 upgrade: - `databases().list` takes `search` and `batch` filters now, so the three call sites pass `None, None`. Behaviour is unchanged: the client omits absent query params and the server applies no filter without them. - `TableInfo` gained required `partition_by` / `sorted_by` fields, so the information-schema mocks in the databases and indexes tests declare them. Without this the mocks fail to deserialize and the test process aborts rather than failing an assertion. --- Cargo.lock | 4 +- Cargo.toml | 2 +- skills/hotdata/SKILL.md | 2 + src/commands/databases.rs | 131 ++++++++++++++++++++++++++++++++++++-- src/commands/indexes.rs | 12 ++-- src/main.rs | 3 + tests/databases_cli.rs | 16 +++++ 7 files changed, 154 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index aea750f1..0820f95c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1197,9 +1197,9 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "hotdata" -version = "0.10.0" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0385ea667c306c33223aba8e7765261f455284f9a549c067a09e6c1742a732c7" +checksum = "001e3ffff61f7e741cb9cdd33637dd87d552f72c0582d776374eff00540b9b45" dependencies = [ "arrow-array", "arrow-ipc", diff --git a/Cargo.toml b/Cargo.toml index 7b1e342b..3132e979 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 diff --git a/skills/hotdata/SKILL.md b/skills/hotdata/SKILL.md index 12c5996e..fdda7236 100644 --- a/skills/hotdata/SKILL.md +++ b/skills/hotdata/SKILL.md @@ -95,6 +95,7 @@ Returns workspaces with `public_id`, `name`, `active`, `favorite`, `provision_st ``` hotdata databases list [--workspace-id ] [--output table|json|yaml] +hotdata databases count [--workspace-id ] [--output table|json|yaml] hotdata databases create [--name ] [--catalog ] [--table ...] [--schema public] [--expires-at ] [--workspace-id ] [--output table|json|yaml] hotdata databases fork [] [--name ] [--expires-at ] [--workspace-id ] [--output table|json|yaml] hotdata databases set @@ -118,6 +119,7 @@ hotdata databases tables delete
[--database ] [--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 .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 `` to fork another. `--name` defaults to `-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. diff --git a/src/commands/databases.rs b/src/commands/databases.rs index 6097df2b..388d5a81 100644 --- a/src/commands/databases.rs +++ b/src/commands/databases.rs @@ -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, 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, 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 { + 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(); diff --git a/src/commands/indexes.rs b/src/commands/indexes.rs index bd4223ba..7bf45277 100644 --- a/src/commands/indexes.rs +++ b/src/commands/indexes.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(); diff --git a/src/main.rs b/src/main.rs index 95b4fc7a..cdee9d50 100644 --- a/src/main.rs +++ b/src/main.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) } diff --git a/tests/databases_cli.rs b/tests/databases_cli.rs index 3b796349..c325cfe4 100644 --- a/tests/databases_cli.rs +++ b/tests/databases_cli.rs @@ -14,6 +14,7 @@ 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")); @@ -21,6 +22,21 @@ fn databases_help_lists_subcommands() { 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()