fix: scope results and query runs to a database

Author: zfarrell

CHANGELOG.md

@@ -10,6 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - feat: support async table loads and append mode
+- **Breaking:** results and query runs are now scoped to a database via the
+  required `X-Database-Id` header. The ergonomic wrappers gain a `database_id`
+  argument to match: `Client::get_result`, `Client::list_results`,
+  `Client::list_query_runs`, `Client::await_result`, `Client::get_result_arrow`,
+  `Client::stream_result_arrow`, `Client::query_to_arrow`, and the
+  `results()` / `query_runs()` resource handles. `Client::query`'s truncation
+  auto-follow now forwards the query's database scope (the `X-Database-Id`
+  header, or the request-body `database_id` when no header is set) to the
+  follow-up result and query-run fetches.
 
 ## [0.7.0] - 2026-06-30
 

README.md

@@ -64,16 +64,22 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         // .base_url("https://api.hotdata.dev")  // optional
         .build()?;
 
-    // Submit a query. `query` transparently retries HTTP 429 (server overload)
+    // Queries, results, and query runs are scoped to a database (the required
+    // `X-Database-Id` header), so pick one first.
+    let database_id = "your_database_id";
+
+    // Submit a query. `query_in` transparently retries HTTP 429 (server overload)
     // and auto-follows a truncated large result to its full row set. Rows come
     // back inline, plus a result_id persisted for later retrieval.
     let response = client
-        .query(QueryRequest::new("SELECT 1 AS n".to_string()))
+        .query_in(QueryRequest::new("SELECT 1 AS n".to_string()), database_id)
         .await?;
 
     if let Some(result_id) = response.result_id.flatten() {
         // Poll the persisted result to `ready` without hand-rolling a loop.
-        let result = client.await_result(&result_id, PollConfig::default()).await?;
+        let result = client
+            .await_result(&result_id, database_id, PollConfig::default())
+            .await?;
         println!("result {} is {}", result.result_id, result.status);
     }
 
@@ -91,7 +97,8 @@ ergonomic, workspace-scoped handles so you never pass a `Configuration` around:
 let connections = client.connections().list().await?;
 let connection  = client.connections().get(&connections.connections[0].id).await?;
 let secrets  = client.secrets().list().await?;
-let runs     = client.query_runs().list(Some(50), None, None, None).await?;
+// Query runs are database-scoped: pass the database id first.
+let runs     = client.query_runs().list(database_id, Some(50), None, None, None).await?;
 ```
 
 Handles exist for every resource — `connections`, `connection_types`,
@@ -101,7 +108,9 @@ Handles exist for every resource — `connections`, `connection_types`,
 operations also have flat shortcuts directly on `Client` (`query` — with
 `query_in` to scope to a database, `query_preview` to skip auto-follow, and
 `query_with` for a per-call `QueryConfig` — plus `get_result`, `list_results`,
-`list_query_runs`, `list_workspaces`).
+`list_query_runs`, `list_workspaces`). The result and query-run operations take
+a `database_id` argument — they are scoped to a database via the required
+`X-Database-Id` header.
 
 For anything not yet wrapped, the full generated surface is one call away via
 `client.configuration()`:
@@ -120,12 +129,12 @@ them with the typed [`ResultStatus`] / [`QueryRunStatus`] enums via the
 ```rust
 use hotdata::prelude::*;
 
-let result = client.await_result(&result_id, PollConfig::default()).await?;
+let result = client.await_result(&result_id, database_id, PollConfig::default()).await?;
 if result.result_status().is_ready() {
     // ...
 }
 
-let run = client.query_runs().get(&query_run_id).await?;
+let run = client.query_runs().get(&query_run_id, database_id).await?;
 if run.run_status().is_terminal() { /* ... */ }
 ```
 
@@ -183,16 +192,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .workspace_id("your_workspace_id")
         .build()?;
 
+    // Results are database-scoped (the required `X-Database-Id` header).
+    let database_id = "your_database_id";
+
     // Buffered: decodes every batch into a Vec<RecordBatch>.
-    let result = client.get_result_arrow(&result_id, None, None).await?;
+    let result = client.get_result_arrow(&result_id, database_id, None, None).await?;
     println!("schema: {:?}", result.schema);
     println!("total rows: {:?}", result.total_row_count);
     for batch in &result.batches {
         // work with each arrow_array::RecordBatch
     }
 
     // Streaming: yields batches lazily without holding them all at once.
-    let mut stream = client.stream_result_arrow(&result_id, None, None).await?;
+    let mut stream = client.stream_result_arrow(&result_id, database_id, None, None).await?;
     for batch in stream.by_ref() {
         let batch = batch?;
         // ...
@@ -202,7 +214,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 }
 ```
 
-Both methods accept `offset` and `limit` for pagination, and both honor the transparent JWT exchange. They return `ArrowError::NotReady` if the result is still pending or processing — poll `client.get_result(result_id)` until its status is `ready` first. `ArrowResult` also surfaces the `X-Total-Row-Count` header (`total_row_count`) and the `rel="next"` pagination `Link` (`next_link`).
+Both methods accept `offset` and `limit` for pagination, and both honor the transparent JWT exchange. They return `ArrowError::NotReady` if the result is still pending or processing — poll `client.get_result(result_id, database_id)` until its status is `ready` first. `ArrowResult` also surfaces the `X-Total-Row-Count` header (`total_row_count`) and the `rel="next"` pagination `Link` (`next_link`).
 
 To run a query and get its result as Arrow in a single call — submit, await
 `ready`, and decode — use `query_to_arrow`:
@@ -211,6 +223,7 @@ To run a query and get its result as Arrow in a single call — submit, await
 let arrow = client
     .query_to_arrow(
         QueryRequest::new("SELECT * FROM big_table".to_string()),
+        database_id,
         PollConfig::default(),
         None, // offset
         None, // limit

examples/quickstart.rs

@@ -92,18 +92,36 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
 
     // --- 3. Submit a query -----------------------------------------------
     //
+    // Queries, results, and query runs are all scoped to a database via the
+    // `X-Database-Id` header, so resolve one first. Prefer HOTDATA_DATABASE, else
+    // fall back to the first database in the workspace; if there are none, skip
+    // the query portion of the tour.
+    let database_id = match resolve_database_id(&client).await? {
+        Some(id) => id,
+        None => {
+            println!(
+                "No database available (set HOTDATA_DATABASE or create one). \
+                 Skipping the query portion of the tour."
+            );
+            return Ok(());
+        }
+    };
+    println!("Using database {database_id}");
+
     // POST /query returns rows inline *and* a result_id; persistence to the
     // result store then completes asynchronously.
     //
-    // `query` is the enhanced default: it retries HTTP 429 (`OVERLOADED`)
-    // transparently and, if the server truncates a large result, auto-follows it
-    // — paging the full row set into `response.rows` — bounded by the instance
-    // `QueryConfig` (default ceilings: 1M rows / 64 MiB). Exceed a ceiling and you
-    // get `QueryError::Result(ResultError::TooLarge { .. })` instead of an OOM.
+    // `query_in` is the database-scoped form of the enhanced default `query`: it
+    // retries HTTP 429 (`OVERLOADED`) transparently and, if the server truncates
+    // a large result, auto-follows it — paging the full row set into
+    // `response.rows` — bounded by the instance `QueryConfig` (default ceilings:
+    // 1M rows / 64 MiB). Exceed a ceiling and you get
+    // `QueryError::Result(ResultError::TooLarge { .. })` instead of an OOM.
     let response = client
-        .query(QueryRequest::new(
-            "select 1 as id, 'hello' as greeting".to_string(),
-        ))
+        .query_in(
+            QueryRequest::new("select 1 as id, 'hello' as greeting".to_string()),
+            &database_id,
+        )
         .await?;
 
     println!(
@@ -112,11 +130,15 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
     );
 
     // Tuning the auto-follow behavior per call: clone the instance config and
-    // override just what you need. Here we opt out of auto-follow entirely —
-    // `query_preview` is the one-call shortcut for "give me only the inline
-    // preview" (handy when you'll page the full result yourself, e.g. as Arrow).
+    // override just what you need. Here we opt out of auto-follow entirely, while
+    // still scoping the query to the database via `query_with`.
+    let preview_config = client.query_config().clone().with_auto_follow(false);
     let preview = client
-        .query_preview(QueryRequest::new("select 1 as id".to_string()))
+        .query_with(
+            QueryRequest::new("select 1 as id".to_string()),
+            Some(&database_id),
+            &preview_config,
+        )
         .await?;
     println!(
         "Preview: {} row(s){}",
@@ -151,24 +173,44 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
     // timeout polled every second.
     println!("Awaiting result {result_id}...");
     let ready = client
-        .await_result(&result_id, PollConfig::default())
+        .await_result(&result_id, &database_id, PollConfig::default())
         .await?;
     println!("Result status: {}", ready.status);
 
     // --- 5. Fetch the result as Arrow (feature-gated) --------------------
-    fetch_arrow(&client, &result_id).await?;
+    fetch_arrow(&client, &database_id, &result_id).await?;
 
     // --- 6. One-call query -> Arrow (feature-gated) ----------------------
-    one_shot_arrow(&client).await?;
+    one_shot_arrow(&client, &database_id).await?;
 
     Ok(())
 }
 
+/// Resolve a database to scope the query tour to: `HOTDATA_DATABASE` if set,
+/// otherwise the first database visible in the workspace (`None` if there are
+/// none).
+async fn resolve_database_id(client: &Client) -> Result<Option<String>, Box<dyn std::error::Error>> {
+    if let Ok(id) = std::env::var("HOTDATA_DATABASE") {
+        if !id.is_empty() {
+            return Ok(Some(id));
+        }
+    }
+    let databases = client.databases().list().await?;
+    Ok(databases.databases.into_iter().next().map(|db| db.id))
+}
+
 /// Fetch an already-ready result as Arrow record batches.
 #[cfg(feature = "arrow")]
-async fn fetch_arrow(client: &Client, result_id: &str) -> Result<(), Box<dyn std::error::Error>> {
+async fn fetch_arrow(
+    client: &Client,
+    database_id: &str,
+    result_id: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
     println!("Fetching result {result_id} as Arrow...");
-    match client.get_result_arrow(result_id, None, None).await {
+    match client
+        .get_result_arrow(result_id, database_id, None, None)
+        .await
+    {
         Ok(arrow) => print_arrow(&arrow),
         // The Arrow error enum maps the result endpoint's status codes to named
         // variants, so callers react without string-matching on HTTP codes.
@@ -183,7 +225,10 @@ async fn fetch_arrow(client: &Client, result_id: &str) -> Result<(), Box<dyn std
 /// Submit a fresh query and get its result as Arrow in a single call —
 /// `query_to_arrow` runs the query, awaits `ready`, and decodes the stream.
 #[cfg(feature = "arrow")]
-async fn one_shot_arrow(client: &Client) -> Result<(), Box<dyn std::error::Error>> {
+async fn one_shot_arrow(
+    client: &Client,
+    database_id: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
     use std::time::Duration;
 
     println!("One-call query_to_arrow...");
@@ -194,6 +239,7 @@ async fn one_shot_arrow(client: &Client) -> Result<(), Box<dyn std::error::Error
     let arrow = client
         .query_to_arrow(
             QueryRequest::new("select 42 as answer".to_string()),
+            database_id,
             poll,
             None,
             None,

src/arrow.rs

@@ -251,11 +251,12 @@ impl Iterator for ArrowBatchStream {
 pub async fn get_result_arrow(
     configuration: &Configuration,
     id: &str,
+    x_database_id: &str,
     offset: Option<i64>,
     limit: Option<i64>,
 ) -> Result<ArrowResult, ArrowError> {
     let (bytes, total_row_count, next_link) =
-        fetch_arrow_bytes(configuration, id, offset, limit).await?;
+        fetch_arrow_bytes(configuration, id, x_database_id, offset, limit).await?;
     let reader = StreamReader::try_new(Cursor::new(bytes), None)?;
     let schema = reader.schema();
     let mut batches = Vec::new();
@@ -283,11 +284,12 @@ pub async fn get_result_arrow(
 pub async fn stream_result_arrow(
     configuration: &Configuration,
     id: &str,
+    x_database_id: &str,
     offset: Option<i64>,
     limit: Option<i64>,
 ) -> Result<ArrowBatchStream, ArrowError> {
     let (bytes, total_row_count, next_link) =
-        fetch_arrow_bytes(configuration, id, offset, limit).await?;
+        fetch_arrow_bytes(configuration, id, x_database_id, offset, limit).await?;
     let reader = StreamReader::try_new(Cursor::new(bytes), None)?;
     Ok(ArrowBatchStream {
         reader,
@@ -326,6 +328,7 @@ fn apply_apikey_headers(
 async fn fetch_arrow_bytes(
     configuration: &Configuration,
     id: &str,
+    x_database_id: &str,
     offset: Option<i64>,
     limit: Option<i64>,
 ) -> Result<(Bytes, Option<i64>, Option<String>), ArrowError> {
@@ -336,6 +339,10 @@ async fn fetch_arrow_bytes(
     );
     let mut req_builder = configuration.client.request(reqwest::Method::GET, &uri_str);
 
+    // `GET /v1/results/{id}` is scoped to a database: the `X-Database-Id` header
+    // is required (mirrors the generated `get_result`).
+    req_builder = req_builder.header("X-Database-Id", x_database_id.to_string());
+
     // format=arrow takes precedence over the Accept header server-side, but we
     // send both to match the generated client and be explicit on the wire.
     req_builder = req_builder.query(&[("format", "arrow")]);
@@ -665,6 +672,38 @@ mod tests {
         assert_eq!(result.total_row_count, Some(5));
     }
 
+    /// Regression (database-scoped results): the Arrow path must send the
+    /// required `X-Database-Id` header on `GET /v1/results/{id}`. The mock
+    /// matches only when the header is present, so a missing header 404s and the
+    /// fetch fails.
+    #[tokio::test]
+    async fn fetch_arrow_sends_database_id_header() {
+        use wiremock::matchers::{header, method, path, query_param};
+        use wiremock::{Mock, MockServer, ResponseTemplate};
+
+        let (ipc, _schema) = make_ipc_stream();
+        let server = MockServer::start().await;
+        Mock::given(method("GET"))
+            .and(path("/v1/results/res_1"))
+            .and(query_param("format", "arrow"))
+            .and(header("X-Database-Id", "db_x"))
+            .respond_with(
+                ResponseTemplate::new(200)
+                    .insert_header("content-type", ARROW_STREAM_MEDIA_TYPE)
+                    .set_body_bytes(ipc),
+            )
+            .mount(&server)
+            .await;
+
+        let mut configuration = Configuration::new();
+        configuration.base_path = server.uri();
+
+        let result = get_result_arrow(&configuration, "res_1", "db_x", None, None)
+            .await
+            .expect("arrow fetch should forward X-Database-Id and succeed");
+        assert_eq!(result.num_rows(), 5);
+    }
+
     #[test]
     fn empty_stream_decodes_to_zero_batches() {
         let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));