From 33c11131e79bd07ff251d5ead6c4801d83ba0ca9 Mon Sep 17 00:00:00 2001 From: "hotdata-automation[bot]" <267177015+hotdata-automation[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 03:42:05 +0000 Subject: [PATCH 1/2] feat: support async table loads and append mode --- CHANGELOG.md | 4 ++++ docs/AsyncQueryResponse.md | 2 +- docs/ConnectionsApi.md | 2 +- docs/CreateIndexRequest.md | 1 + docs/DatabasesApi.md | 2 +- docs/JobResult.md | 1 + docs/LoadManagedTableRequest.md | 4 +++- docs/LoadManagedTableResponse.md | 2 +- docs/QueryRequest.md | 2 +- docs/QueryRunsApi.md | 10 +++++++--- docs/RefreshRequest.md | 1 + docs/RefreshResponse.md | 1 - docs/ResultsApi.md | 8 ++++++-- src/apis/connections_api.rs | 2 +- src/apis/databases_api.rs | 2 +- src/apis/query_runs_api.rs | 12 +++++++++++- src/apis/refresh_api.rs | 1 + src/apis/results_api.rs | 9 +++++++++ src/models/async_query_response.rs | 6 +++--- src/models/create_index_request.rs | 9 +++++++++ src/models/job_result.rs | 1 + src/models/load_managed_table_request.rs | 19 ++++++++++++++++--- src/models/load_managed_table_response.rs | 6 +++--- src/models/query_request.rs | 2 +- src/models/refresh_request.rs | 9 +++++++++ src/models/refresh_response.rs | 1 - 26 files changed, 93 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 896747f..fc10e6d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- feat: support async table loads and append mode + ## [0.7.0] - 2026-06-30 ### Added diff --git a/docs/AsyncQueryResponse.md b/docs/AsyncQueryResponse.md index 2565b60..237ee2b 100644 --- a/docs/AsyncQueryResponse.md +++ b/docs/AsyncQueryResponse.md @@ -7,7 +7,7 @@ Name | Type | Description | Notes **query_run_id** | **String** | Unique identifier for the query run. | **reason** | Option<**String**> | Human-readable reason why the query went async (e.g., caching tables for the first time). | [optional] **status** | **String** | Current status of the query run. | -**status_url** | **String** | URL to poll for query run status. | +**status_url** | **String** | URL to poll for query run status. Requires the same `X-Database-Id` header used to submit the query. | [[Back to Model list]](../README.md#documentation-for-models) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to README]](../README.md) diff --git a/docs/ConnectionsApi.md b/docs/ConnectionsApi.md index 6c3eada..3ecfb9d 100644 --- a/docs/ConnectionsApi.md +++ b/docs/ConnectionsApi.md @@ -298,7 +298,7 @@ This endpoint does not need any parameter. > models::LoadManagedTableResponse load_managed_table(connection_id, schema, table, load_managed_table_request) Load managed table from upload -Publish a previously-uploaded file as the new contents of a managed table. CSV, JSON, and Parquet uploads are supported; the format is auto-detected from the upload's `Content-Type` and file contents, or set explicitly via the `format` field. Only `mode = \"replace\"` is supported. Concurrent loads against the same upload return 409. +Publish a previously-uploaded file as the new contents of a managed table. CSV, JSON, and Parquet uploads are supported; the format is auto-detected from the upload's `Content-Type` and file contents, or set explicitly via the `format` field. If the target table (or its schema) has not been declared yet, it is created automatically as part of the load — declaring tables up front is optional. `mode` selects how the upload is applied: `replace` overwrites the table's contents, `append` inserts the uploaded rows on top of the existing data. Concurrent loads against the same upload return 409. Set `async` to run the load in the background and get back a job ID to poll; add `async_after_ms` to wait briefly for it to finish before falling back to a job ID. ### Parameters diff --git a/docs/CreateIndexRequest.md b/docs/CreateIndexRequest.md index fb9e955..bd2c3c0 100644 --- a/docs/CreateIndexRequest.md +++ b/docs/CreateIndexRequest.md @@ -5,6 +5,7 @@ Name | Type | Description | Notes ------------ | ------------- | ------------- | ------------- **r#async** | Option<**bool**> | When true, create the index as a background job and return a job ID for polling. | [optional] +**async_after_ms** | Option<**i32**> | If set (requires `async` = true), wait up to this many milliseconds for the index build to finish: if it completes in time the index is returned (201), otherwise a 202 with a job ID to poll. Must be between 1000 and the server maximum; a value out of that range, or set without `async` = true, is rejected with 400. | [optional] **columns** | **Vec** | Columns to index. Required for all index types. | **description** | Option<**String**> | User-facing description of the embedding (e.g., \"product descriptions\"). | [optional] **dimensions** | Option<**i32**> | Output vector dimensions. Some models support multiple dimension sizes (e.g., OpenAI text-embedding-3-small supports 512 or 1536). If omitted, the model's default dimensions are used | [optional] diff --git a/docs/DatabasesApi.md b/docs/DatabasesApi.md index 0408768..a7e9498 100644 --- a/docs/DatabasesApi.md +++ b/docs/DatabasesApi.md @@ -259,7 +259,7 @@ This endpoint does not need any parameter. > models::LoadManagedTableResponse load_database_table(database_id, schema, table, load_managed_table_request) Load database table from upload -Publish a previously-uploaded file as the new contents of a table on the database's default catalog. The database-scoped equivalent of the connection-scoped managed-table load — addressed by `database_id`, so no `default_connection_id` is needed. CSV, JSON, and Parquet uploads are supported; the format is auto-detected or set via `format`. Only `mode = \"replace\"` is supported. Concurrent loads against the same upload return 409. +Publish a previously-uploaded file as the new contents of a table on the database's default catalog. The database-scoped equivalent of the connection-scoped managed-table load — addressed by `database_id`, so no `default_connection_id` is needed. CSV, JSON, and Parquet uploads are supported; the format is auto-detected or set via `format`. If the target table (or its schema) has not been declared yet, it is created automatically as part of the load — declaring tables up front is optional. `mode` selects how the upload is applied: `replace` overwrites the table's contents, `append` inserts the uploaded rows on top of the existing data. Concurrent loads against the same upload return 409. Set `async` to run the load in the background and get back a job ID to poll; add `async_after_ms` to wait briefly for it to finish before falling back to a job ID. ### Parameters diff --git a/docs/JobResult.md b/docs/JobResult.md index 06636f0..ceebeef 100644 --- a/docs/JobResult.md +++ b/docs/JobResult.md @@ -6,6 +6,7 @@ |---- | -----| | ConnectionRefreshResult | Job-specific result payload. The shape depends on the job type. Null while the job is pending or running. | | IndexInfoResponse | Job-specific result payload. The shape depends on the job type. Null while the job is pending or running. | +| LoadManagedTableResponse | Job-specific result payload. The shape depends on the job type. Null while the job is pending or running. | | TableRefreshResult | Job-specific result payload. The shape depends on the job type. Null while the job is pending or running. | [[Back to Model list]](../README.md#documentation-for-models) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to README]](../README.md) diff --git a/docs/LoadManagedTableRequest.md b/docs/LoadManagedTableRequest.md index 3806ca3..624184f 100644 --- a/docs/LoadManagedTableRequest.md +++ b/docs/LoadManagedTableRequest.md @@ -4,8 +4,10 @@ Name | Type | Description | Notes ------------ | ------------- | ------------- | ------------- +**r#async** | Option<**bool**> | When true, run the load as a background job and return a job ID to poll instead of blocking until it finishes. Recommended for large uploads, which can take longer than an HTTP request should stay open. | [optional] +**async_after_ms** | Option<**i32**> | If set (requires `async` = true), wait up to this many milliseconds for the load to finish: if it completes in time the full result is returned (200), otherwise a 202 with a job ID to poll. Must be between 1000 and the server maximum; a value out of that range, or set without `async` = true, is rejected with 400. | [optional] **format** | Option<**String**> | File format of the upload: `\"csv\"`, `\"json\"`, or `\"parquet\"`. Optional — when omitted, the format is auto-detected from the upload's `Content-Type` and, failing that, from the file contents. Provide it explicitly to override detection or when the contents are ambiguous. `\"json\"` expects newline-delimited JSON (one object per line), not a JSON array. | [optional] -**mode** | **String** | Load mode. Only `\"replace\"` is supported in this release. | +**mode** | **String** | How the upload is applied: `\"replace\"` overwrites the table's contents, `\"append\"` inserts the uploaded rows on top of the existing data. | **upload_id** | **String** | ID of a previously-staged upload (see `POST /v1/files`). The upload is claimed atomically; concurrent loads against the same `upload_id` return 409. | [[Back to Model list]](../README.md#documentation-for-models) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to README]](../README.md) diff --git a/docs/LoadManagedTableResponse.md b/docs/LoadManagedTableResponse.md index 0365d6a..9463f6a 100644 --- a/docs/LoadManagedTableResponse.md +++ b/docs/LoadManagedTableResponse.md @@ -6,7 +6,7 @@ Name | Type | Description | Notes ------------ | ------------- | ------------- | ------------- **arrow_schema_json** | **String** | Schema of the loaded table, as JSON. | **connection_id** | **String** | | -**row_count** | **i64** | Total rows in the published parquet file. | +**row_count** | **i64** | Total number of rows in the table after the load. | **schema_name** | **String** | | **table_name** | **String** | | diff --git a/docs/QueryRequest.md b/docs/QueryRequest.md index d87bb19..8d64687 100644 --- a/docs/QueryRequest.md +++ b/docs/QueryRequest.md @@ -5,7 +5,7 @@ Name | Type | Description | Notes ------------ | ------------- | ------------- | ------------- **r#async** | Option<**bool**> | When true, execute the query asynchronously and return a query run ID for polling via GET /query-runs/{id}. The query results can be retrieved via GET /results/{id} once the query run status is \"succeeded\". | [optional] -**async_after_ms** | Option<**i32**> | If set with async=true, wait up to this many milliseconds for the query to complete synchronously before returning an async response. Minimum 1000ms. Ignored if async is false. | [optional] +**async_after_ms** | Option<**i32**> | If set (requires `async` = true), first attempt the query synchronously and wait up to this many milliseconds: if it finishes in time the full result is returned, otherwise an async response (a run id to poll) is returned. Must be at least 1000 and at most the server's configured maximum; a value out of that range, or set without `async` = true, is rejected with 400. | [optional] **database_id** | Option<**String**> | Database to scope the query to (its id). Alternative to the `X-Database-Id` header — exactly one source must be provided. If both this field and the header are set and they disagree, the request is rejected with a 400. | [optional] **default_catalog** | Option<**String**> | Catalog that unqualified table references resolve against within the query's database scope. Must name a catalog visible in the database (`default`, an attached catalog alias, or a system catalog). Defaults to `default` when omitted. | [optional] **default_schema** | Option<**String**> | Schema that unqualified table references resolve against within the query's database scope. Defaults to `main` when omitted. Existence is not validated up front — an unknown schema surfaces as a \"table not found\" error at planning time. | [optional] diff --git a/docs/QueryRunsApi.md b/docs/QueryRunsApi.md index aaf74ce..0685804 100644 --- a/docs/QueryRunsApi.md +++ b/docs/QueryRunsApi.md @@ -11,10 +11,10 @@ Method | HTTP request | Description ## get_query_run -> models::QueryRunInfo get_query_run(id) +> models::QueryRunInfo get_query_run(id, x_database_id) Get query run -Get the status and details of a specific query run by ID. +Get the status and details of a specific query run by ID, scoped to the database named by the required X-Database-Id header. ### Parameters @@ -22,6 +22,7 @@ Get the status and details of a specific query run by ID. Name | Type | Description | Required | Notes ------------- | ------------- | ------------- | ------------- | ------------- **id** | **String** | Query run ID | [required] | +**x_database_id** | **String** | Database the query run belongs to (required) | [required] | ### Return type @@ -41,14 +42,17 @@ Name | Type | Description | Required | Notes ## list_query_runs -> models::ListQueryRunsResponse list_query_runs(limit, cursor, status, saved_query_id) +> models::ListQueryRunsResponse list_query_runs(x_database_id, limit, cursor, status, saved_query_id) List query runs +List query runs for the database named by the required X-Database-Id header. + ### Parameters Name | Type | Description | Required | Notes ------------- | ------------- | ------------- | ------------- | ------------- +**x_database_id** | **String** | Database to scope the query runs to (required) | [required] | **limit** | Option<**i32**> | Maximum number of results | | **cursor** | Option<**String**> | Pagination cursor | | **status** | Option<**String**> | Filter by status (comma-separated, e.g. status=running,failed) | | diff --git a/docs/RefreshRequest.md b/docs/RefreshRequest.md index d2fd148..e5fabf4 100644 --- a/docs/RefreshRequest.md +++ b/docs/RefreshRequest.md @@ -5,6 +5,7 @@ Name | Type | Description | Notes ------------ | ------------- | ------------- | ------------- **r#async** | Option<**bool**> | When true, submit the refresh as a background job and return immediately with a job ID for status polling. Only supported for data refresh operations. | [optional] +**async_after_ms** | Option<**i32**> | If set (requires `async` = true), wait up to this many milliseconds for the refresh to finish: if it completes in time the full result is returned, otherwise a `202` with a job ID to poll. Must be between 1000 and the server maximum; a value out of that range, or set without `async` = true, is rejected with 400. Only applies to data refresh. | [optional] **connection_id** | Option<**String**> | | [optional] **data** | Option<**bool**> | | [optional] **include_uncached** | Option<**bool**> | Controls whether uncached tables are included in connection-wide data refresh. - `false` (default): Only refresh tables that already have cached data. This is the common case for keeping existing data up-to-date. - `true`: Also sync tables that haven't been cached yet, essentially performing an initial sync for any new tables discovered since the connection was created. This field only applies to connection-wide data refresh (when `data=true` and `table_name` is not specified). It has no effect on single-table refresh or schema refresh operations. | [optional] diff --git a/docs/RefreshResponse.md b/docs/RefreshResponse.md index 9a518ba..6ad2ad1 100644 --- a/docs/RefreshResponse.md +++ b/docs/RefreshResponse.md @@ -6,7 +6,6 @@ |---- | -----| | ConnectionRefreshResult | Unified response type for refresh operations | | SchemaRefreshResult | Unified response type for refresh operations | -| SubmitJobResponse | Unified response type for refresh operations | | TableRefreshResult | Unified response type for refresh operations | [[Back to Model list]](../README.md#documentation-for-models) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to README]](../README.md) diff --git a/docs/ResultsApi.md b/docs/ResultsApi.md index 1dc4685..aadf961 100644 --- a/docs/ResultsApi.md +++ b/docs/ResultsApi.md @@ -11,7 +11,7 @@ Method | HTTP request | Description ## get_result -> models::GetResultResponse get_result(id, offset, limit, format) +> models::GetResultResponse get_result(id, x_database_id, offset, limit, format) Get result Retrieve a persisted query result by ID. The response format for the `ready` state is selected by `Accept` header or `?format=` query param; non-ready states use the same status codes and JSON body shape regardless of format. | Result status | Status × body | |-----------------------|------------------------------------------------------------------------------| | `ready` + JSON | 200 `application/json` — `GetResultResponse` with `columns`, `rows`, etc. | | `ready` + Arrow | 200 `application/vnd.apache.arrow.stream` — schema, RecordBatches, EOS | | `ready` + CSV | 200 `text/csv; charset=utf-8` — single header row, streamed batch-by-batch | | `ready` + Markdown | 200 `text/markdown; charset=utf-8` — GitHub-flavored pipe table, streamed | | `ready` + Parquet | 200 `application/vnd.apache.parquet` — raw parquet bytes (no conversion) | | `pending`/`processing`| 202 `application/json` `{status, result_id}` + `Retry-After` | | `failed` | 409 `application/json` `{status, result_id, error_message}` | | not found | 404 `application/json` (`ApiErrorResponse`) | `?format=` accepts `arrow`, `json`, `csv`, `md`, `parquet` and takes precedence over `Accept`. `markdown` is accepted as a runtime alias for `md`. Use `?offset=N&limit=M` to slice the result; `offset` defaults to 0 and `limit` is unbounded by default. Both must be non-negative; invalid values return 400. When a finite `limit` doesn't reach the end of the result, a `Link` header with `rel=\"next\"` points at the following page. `?offset`/`?limit` are ignored for `format=parquet` since that path returns the underlying file unchanged. Ready responses (Arrow, CSV, Markdown, JSON) carry `X-Total-Row-Count` (the full result row count, independent of offset/limit). Responses are streamed end-to-end, so a client can disconnect at any time and the server stops reading. IEEE special floats (`±Inf`, `NaN`) have no canonical JSON representation. For cross-format consistency the JSON, CSV, and Markdown paths emit them as `null` / empty cells, and JSON `nullable[]` is widened to match. The Arrow IPC and Parquet bodies are binary round-trip formats and preserve the raw IEEE values; callers cross-checking a result across CSV and Parquet should not byte-compare those slots. @@ -22,6 +22,7 @@ Retrieve a persisted query result by ID. The response format for the `ready` sta Name | Type | Description | Required | Notes ------------- | ------------- | ------------- | ------------- | ------------- **id** | **String** | Result ID | [required] | +**x_database_id** | **String** | Database the result belongs to (required) | [required] | **offset** | Option<**i32**> | Rows to skip (default: 0) | | **limit** | Option<**i32**> | Maximum rows to return (default: unbounded) | | **format** | Option<[**ResultsFormatQuery**](ResultsFormatQuery.md)> | `arrow`, `json`, `csv`, `md`, or `parquet` — overrides the `Accept` header. `markdown` is also accepted at runtime as an alias for `md`. | | @@ -44,14 +45,17 @@ Name | Type | Description | Required | Notes ## list_results -> models::ListResultsResponse list_results(limit, offset) +> models::ListResultsResponse list_results(x_database_id, limit, offset) List results +List stored results for the database named by the required X-Database-Id header. + ### Parameters Name | Type | Description | Required | Notes ------------- | ------------- | ------------- | ------------- | ------------- +**x_database_id** | **String** | Database to scope the results to (required) | [required] | **limit** | Option<**i32**> | Maximum number of results (default: 100, max: 1000) | | **offset** | Option<**i32**> | Pagination offset (default: 0) | | diff --git a/src/apis/connections_api.rs b/src/apis/connections_api.rs index b1f145c..335d1a8 100644 --- a/src/apis/connections_api.rs +++ b/src/apis/connections_api.rs @@ -718,7 +718,7 @@ pub async fn list_connections( } } -/// Publish a previously-uploaded file as the new contents of a managed table. CSV, JSON, and Parquet uploads are supported; the format is auto-detected from the upload's `Content-Type` and file contents, or set explicitly via the `format` field. Only `mode = \"replace\"` is supported. Concurrent loads against the same upload return 409. +/// Publish a previously-uploaded file as the new contents of a managed table. CSV, JSON, and Parquet uploads are supported; the format is auto-detected from the upload's `Content-Type` and file contents, or set explicitly via the `format` field. If the target table (or its schema) has not been declared yet, it is created automatically as part of the load — declaring tables up front is optional. `mode` selects how the upload is applied: `replace` overwrites the table's contents, `append` inserts the uploaded rows on top of the existing data. Concurrent loads against the same upload return 409. Set `async` to run the load in the background and get back a job ID to poll; add `async_after_ms` to wait briefly for it to finish before falling back to a job ID. pub async fn load_managed_table( configuration: &configuration::Configuration, connection_id: &str, diff --git a/src/apis/databases_api.rs b/src/apis/databases_api.rs index 99b14ef..eeb7118 100644 --- a/src/apis/databases_api.rs +++ b/src/apis/databases_api.rs @@ -608,7 +608,7 @@ pub async fn list_databases( } } -/// Publish a previously-uploaded file as the new contents of a table on the database's default catalog. The database-scoped equivalent of the connection-scoped managed-table load — addressed by `database_id`, so no `default_connection_id` is needed. CSV, JSON, and Parquet uploads are supported; the format is auto-detected or set via `format`. Only `mode = \"replace\"` is supported. Concurrent loads against the same upload return 409. +/// Publish a previously-uploaded file as the new contents of a table on the database's default catalog. The database-scoped equivalent of the connection-scoped managed-table load — addressed by `database_id`, so no `default_connection_id` is needed. CSV, JSON, and Parquet uploads are supported; the format is auto-detected or set via `format`. If the target table (or its schema) has not been declared yet, it is created automatically as part of the load — declaring tables up front is optional. `mode` selects how the upload is applied: `replace` overwrites the table's contents, `append` inserts the uploaded rows on top of the existing data. Concurrent loads against the same upload return 409. Set `async` to run the load in the background and get back a job ID to poll; add `async_after_ms` to wait briefly for it to finish before falling back to a job ID. pub async fn load_database_table( configuration: &configuration::Configuration, database_id: &str, diff --git a/src/apis/query_runs_api.rs b/src/apis/query_runs_api.rs index 98e542b..b953f35 100644 --- a/src/apis/query_runs_api.rs +++ b/src/apis/query_runs_api.rs @@ -17,6 +17,7 @@ use serde::{de::Error as _, Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum GetQueryRunError { + Status400(models::ApiErrorResponse), Status404(models::ApiErrorResponse), UnknownValue(serde_json::Value), } @@ -25,16 +26,20 @@ pub enum GetQueryRunError { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum ListQueryRunsError { + Status400(models::ApiErrorResponse), + Status404(models::ApiErrorResponse), UnknownValue(serde_json::Value), } -/// Get the status and details of a specific query run by ID. +/// Get the status and details of a specific query run by ID, scoped to the database named by the required X-Database-Id header. pub async fn get_query_run( configuration: &configuration::Configuration, id: &str, + x_database_id: &str, ) -> Result> { // add a prefix to parameters to efficiently prevent name collisions let p_path_id = id; + let p_header_x_database_id = x_database_id; let uri_str = format!( "{}/v1/query-runs/{id}", @@ -46,6 +51,7 @@ pub async fn get_query_run( if let Some(ref user_agent) = configuration.user_agent { req_builder = req_builder.header(reqwest::header::USER_AGENT, user_agent.clone()); } + req_builder = req_builder.header("X-Database-Id", p_header_x_database_id.to_string()); if let Some(apikey) = configuration.api_keys.get("X-Workspace-Id") { let key = apikey.key.clone(); let value = match apikey.prefix { @@ -95,14 +101,17 @@ pub async fn get_query_run( } } +/// List query runs for the database named by the required X-Database-Id header. pub async fn list_query_runs( configuration: &configuration::Configuration, + x_database_id: &str, limit: Option, cursor: Option<&str>, status: Option<&str>, saved_query_id: Option<&str>, ) -> Result> { // add a prefix to parameters to efficiently prevent name collisions + let p_header_x_database_id = x_database_id; let p_query_limit = limit; let p_query_cursor = cursor; let p_query_status = status; @@ -126,6 +135,7 @@ pub async fn list_query_runs( if let Some(ref user_agent) = configuration.user_agent { req_builder = req_builder.header(reqwest::header::USER_AGENT, user_agent.clone()); } + req_builder = req_builder.header("X-Database-Id", p_header_x_database_id.to_string()); if let Some(apikey) = configuration.api_keys.get("X-Workspace-Id") { let key = apikey.key.clone(); let value = match apikey.prefix { diff --git a/src/apis/refresh_api.rs b/src/apis/refresh_api.rs index d7a8df1..6b5c038 100644 --- a/src/apis/refresh_api.rs +++ b/src/apis/refresh_api.rs @@ -19,6 +19,7 @@ use serde::{de::Error as _, Deserialize, Serialize}; pub enum RefreshError { Status400(models::ApiErrorResponse), Status404(models::ApiErrorResponse), + Status409(models::ApiErrorResponse), UnknownValue(serde_json::Value), } diff --git a/src/apis/results_api.rs b/src/apis/results_api.rs index 663b946..1442c3d 100644 --- a/src/apis/results_api.rs +++ b/src/apis/results_api.rs @@ -27,6 +27,8 @@ pub enum GetResultError { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum ListResultsError { + Status400(models::ApiErrorResponse), + Status404(models::ApiErrorResponse), UnknownValue(serde_json::Value), } @@ -34,12 +36,14 @@ pub enum ListResultsError { pub async fn get_result( configuration: &configuration::Configuration, id: &str, + x_database_id: &str, offset: Option, limit: Option, format: Option, ) -> Result> { // add a prefix to parameters to efficiently prevent name collisions let p_path_id = id; + let p_header_x_database_id = x_database_id; let p_query_offset = offset; let p_query_limit = limit; let p_query_format = format; @@ -63,6 +67,7 @@ pub async fn get_result( if let Some(ref user_agent) = configuration.user_agent { req_builder = req_builder.header(reqwest::header::USER_AGENT, user_agent.clone()); } + req_builder = req_builder.header("X-Database-Id", p_header_x_database_id.to_string()); if let Some(apikey) = configuration.api_keys.get("X-Workspace-Id") { let key = apikey.key.clone(); let value = match apikey.prefix { @@ -120,12 +125,15 @@ pub async fn get_result( } } +/// List stored results for the database named by the required X-Database-Id header. pub async fn list_results( configuration: &configuration::Configuration, + x_database_id: &str, limit: Option, offset: Option, ) -> Result> { // add a prefix to parameters to efficiently prevent name collisions + let p_header_x_database_id = x_database_id; let p_query_limit = limit; let p_query_offset = offset; @@ -141,6 +149,7 @@ pub async fn list_results( if let Some(ref user_agent) = configuration.user_agent { req_builder = req_builder.header(reqwest::header::USER_AGENT, user_agent.clone()); } + req_builder = req_builder.header("X-Database-Id", p_header_x_database_id.to_string()); if let Some(apikey) = configuration.api_keys.get("X-Workspace-Id") { let key = apikey.key.clone(); let value = match apikey.prefix { diff --git a/src/models/async_query_response.rs b/src/models/async_query_response.rs index bd08eac..387d0c1 100644 --- a/src/models/async_query_response.rs +++ b/src/models/async_query_response.rs @@ -11,7 +11,7 @@ use crate::models; use serde::{Deserialize, Serialize}; -/// AsyncQueryResponse : Response returned when a query is submitted asynchronously (202 Accepted). Poll GET /query-runs/{id} to track progress. Once status is \"succeeded\", retrieve results via GET /results/{result_id}. +/// AsyncQueryResponse : Response returned when a query is submitted asynchronously (202 Accepted). Poll GET /query-runs/{id} to track progress, sending the same `X-Database-Id` header used to submit the query — the endpoint is scoped to that database and returns 400 without it. Once status is \"succeeded\", retrieve results via GET /results/{result_id}. #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] pub struct AsyncQueryResponse { /// Unique identifier for the query run. @@ -28,13 +28,13 @@ pub struct AsyncQueryResponse { /// Current status of the query run. #[serde(rename = "status")] pub status: String, - /// URL to poll for query run status. + /// URL to poll for query run status. Requires the same `X-Database-Id` header used to submit the query. #[serde(rename = "status_url")] pub status_url: String, } impl AsyncQueryResponse { - /// Response returned when a query is submitted asynchronously (202 Accepted). Poll GET /query-runs/{id} to track progress. Once status is \"succeeded\", retrieve results via GET /results/{result_id}. + /// Response returned when a query is submitted asynchronously (202 Accepted). Poll GET /query-runs/{id} to track progress, sending the same `X-Database-Id` header used to submit the query — the endpoint is scoped to that database and returns 400 without it. Once status is \"succeeded\", retrieve results via GET /results/{result_id}. pub fn new(query_run_id: String, status: String, status_url: String) -> AsyncQueryResponse { AsyncQueryResponse { query_run_id, diff --git a/src/models/create_index_request.rs b/src/models/create_index_request.rs index 7ff1285..cb47bf5 100644 --- a/src/models/create_index_request.rs +++ b/src/models/create_index_request.rs @@ -17,6 +17,14 @@ pub struct CreateIndexRequest { /// When true, create the index as a background job and return a job ID for polling. #[serde(rename = "async", skip_serializing_if = "Option::is_none")] pub r#async: Option, + /// If set (requires `async` = true), wait up to this many milliseconds for the index build to finish: if it completes in time the index is returned (201), otherwise a 202 with a job ID to poll. Must be between 1000 and the server maximum; a value out of that range, or set without `async` = true, is rejected with 400. + #[serde( + rename = "async_after_ms", + default, + with = "::serde_with::rust::double_option", + skip_serializing_if = "Option::is_none" + )] + pub async_after_ms: Option>, /// Columns to index. Required for all index types. #[serde(rename = "columns")] pub columns: Vec, @@ -72,6 +80,7 @@ impl CreateIndexRequest { pub fn new(columns: Vec, index_name: String) -> CreateIndexRequest { CreateIndexRequest { r#async: None, + async_after_ms: None, columns, description: None, dimensions: None, diff --git a/src/models/job_result.rs b/src/models/job_result.rs index dfa2d05..934bc3b 100644 --- a/src/models/job_result.rs +++ b/src/models/job_result.rs @@ -19,6 +19,7 @@ pub enum JobResult { TableRefreshResult(Box), ConnectionRefreshResult(Box), IndexInfoResponse(Box), + LoadManagedTableResponse(Box), } impl Default for JobResult { diff --git a/src/models/load_managed_table_request.rs b/src/models/load_managed_table_request.rs index 40c0ba2..e69a2d1 100644 --- a/src/models/load_managed_table_request.rs +++ b/src/models/load_managed_table_request.rs @@ -11,9 +11,20 @@ use crate::models; use serde::{Deserialize, Serialize}; -/// LoadManagedTableRequest : Request body for `POST /v1/connections/{connection_id}/schemas/{schema}/tables/{table}/loads`. Publishes a previously-uploaded file as the new contents of the named managed table. CSV and JSON uploads are converted to columnar storage on load; Parquet uploads are published directly. `mode` is fixed to `\"replace\"` today; the field is kept in the request body so future modes (e.g. append) are an additive change. +/// LoadManagedTableRequest : Request body for the managed-table load endpoints — the connection-scoped `POST /v1/connections/{connection_id}/schemas/{schema}/tables/{table}/loads` and the database-scoped equivalent. Publishes a previously-uploaded file to the named table. CSV and JSON uploads are converted to columnar storage on load; Parquet uploads are published directly. `mode` selects whether the upload replaces the table's contents or is appended on top of them. #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] pub struct LoadManagedTableRequest { + /// When true, run the load as a background job and return a job ID to poll instead of blocking until it finishes. Recommended for large uploads, which can take longer than an HTTP request should stay open. + #[serde(rename = "async", skip_serializing_if = "Option::is_none")] + pub r#async: Option, + /// If set (requires `async` = true), wait up to this many milliseconds for the load to finish: if it completes in time the full result is returned (200), otherwise a 202 with a job ID to poll. Must be between 1000 and the server maximum; a value out of that range, or set without `async` = true, is rejected with 400. + #[serde( + rename = "async_after_ms", + default, + with = "::serde_with::rust::double_option", + skip_serializing_if = "Option::is_none" + )] + pub async_after_ms: Option>, /// File format of the upload: `\"csv\"`, `\"json\"`, or `\"parquet\"`. Optional — when omitted, the format is auto-detected from the upload's `Content-Type` and, failing that, from the file contents. Provide it explicitly to override detection or when the contents are ambiguous. `\"json\"` expects newline-delimited JSON (one object per line), not a JSON array. #[serde( rename = "format", @@ -22,7 +33,7 @@ pub struct LoadManagedTableRequest { skip_serializing_if = "Option::is_none" )] pub format: Option>, - /// Load mode. Only `\"replace\"` is supported in this release. + /// How the upload is applied: `\"replace\"` overwrites the table's contents, `\"append\"` inserts the uploaded rows on top of the existing data. #[serde(rename = "mode")] pub mode: String, /// ID of a previously-staged upload (see `POST /v1/files`). The upload is claimed atomically; concurrent loads against the same `upload_id` return 409. @@ -31,9 +42,11 @@ pub struct LoadManagedTableRequest { } impl LoadManagedTableRequest { - /// Request body for `POST /v1/connections/{connection_id}/schemas/{schema}/tables/{table}/loads`. Publishes a previously-uploaded file as the new contents of the named managed table. CSV and JSON uploads are converted to columnar storage on load; Parquet uploads are published directly. `mode` is fixed to `\"replace\"` today; the field is kept in the request body so future modes (e.g. append) are an additive change. + /// Request body for the managed-table load endpoints — the connection-scoped `POST /v1/connections/{connection_id}/schemas/{schema}/tables/{table}/loads` and the database-scoped equivalent. Publishes a previously-uploaded file to the named table. CSV and JSON uploads are converted to columnar storage on load; Parquet uploads are published directly. `mode` selects whether the upload replaces the table's contents or is appended on top of them. pub fn new(mode: String, upload_id: String) -> LoadManagedTableRequest { LoadManagedTableRequest { + r#async: None, + async_after_ms: None, format: None, mode, upload_id, diff --git a/src/models/load_managed_table_response.rs b/src/models/load_managed_table_response.rs index 4ea337d..b3b643c 100644 --- a/src/models/load_managed_table_response.rs +++ b/src/models/load_managed_table_response.rs @@ -11,7 +11,7 @@ use crate::models; use serde::{Deserialize, Serialize}; -/// LoadManagedTableResponse : Response body for `POST /v1/connections/{connection_id}/schemas/{schema}/tables/{table}/loads`. +/// LoadManagedTableResponse : Result of a managed-table load: row count after the load plus the published table schema. Returned inline (`200`) for a synchronous load, and as the `GET /v1/jobs/{id}` result payload for a completed background load. #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] pub struct LoadManagedTableResponse { /// Schema of the loaded table, as JSON. @@ -19,7 +19,7 @@ pub struct LoadManagedTableResponse { pub arrow_schema_json: String, #[serde(rename = "connection_id")] pub connection_id: String, - /// Total rows in the published parquet file. + /// Total number of rows in the table after the load. #[serde(rename = "row_count")] pub row_count: i64, #[serde(rename = "schema_name")] @@ -29,7 +29,7 @@ pub struct LoadManagedTableResponse { } impl LoadManagedTableResponse { - /// Response body for `POST /v1/connections/{connection_id}/schemas/{schema}/tables/{table}/loads`. + /// Result of a managed-table load: row count after the load plus the published table schema. Returned inline (`200`) for a synchronous load, and as the `GET /v1/jobs/{id}` result payload for a completed background load. pub fn new( arrow_schema_json: String, connection_id: String, diff --git a/src/models/query_request.rs b/src/models/query_request.rs index a521c0b..d50e86c 100644 --- a/src/models/query_request.rs +++ b/src/models/query_request.rs @@ -17,7 +17,7 @@ pub struct QueryRequest { /// When true, execute the query asynchronously and return a query run ID for polling via GET /query-runs/{id}. The query results can be retrieved via GET /results/{id} once the query run status is \"succeeded\". #[serde(rename = "async", skip_serializing_if = "Option::is_none")] pub r#async: Option, - /// If set with async=true, wait up to this many milliseconds for the query to complete synchronously before returning an async response. Minimum 1000ms. Ignored if async is false. + /// If set (requires `async` = true), first attempt the query synchronously and wait up to this many milliseconds: if it finishes in time the full result is returned, otherwise an async response (a run id to poll) is returned. Must be at least 1000 and at most the server's configured maximum; a value out of that range, or set without `async` = true, is rejected with 400. #[serde( rename = "async_after_ms", default, diff --git a/src/models/refresh_request.rs b/src/models/refresh_request.rs index 76333e2..603ee8a 100644 --- a/src/models/refresh_request.rs +++ b/src/models/refresh_request.rs @@ -17,6 +17,14 @@ pub struct RefreshRequest { /// When true, submit the refresh as a background job and return immediately with a job ID for status polling. Only supported for data refresh operations. #[serde(rename = "async", skip_serializing_if = "Option::is_none")] pub r#async: Option, + /// If set (requires `async` = true), wait up to this many milliseconds for the refresh to finish: if it completes in time the full result is returned, otherwise a `202` with a job ID to poll. Must be between 1000 and the server maximum; a value out of that range, or set without `async` = true, is rejected with 400. Only applies to data refresh. + #[serde( + rename = "async_after_ms", + default, + with = "::serde_with::rust::double_option", + skip_serializing_if = "Option::is_none" + )] + pub async_after_ms: Option>, #[serde( rename = "connection_id", default, @@ -50,6 +58,7 @@ impl RefreshRequest { pub fn new() -> RefreshRequest { RefreshRequest { r#async: None, + async_after_ms: None, connection_id: None, data: None, include_uncached: None, diff --git a/src/models/refresh_response.rs b/src/models/refresh_response.rs index fa3157f..1602496 100644 --- a/src/models/refresh_response.rs +++ b/src/models/refresh_response.rs @@ -19,7 +19,6 @@ pub enum RefreshResponse { SchemaRefreshResult(Box), TableRefreshResult(Box), ConnectionRefreshResult(Box), - SubmitJobResponse(Box), } impl Default for RefreshResponse { From 431c19552564494f601b386aa7849e7b42ebaa73 Mon Sep 17 00:00:00 2001 From: Zac Farrell Date: Tue, 7 Jul 2026 08:57:48 -0700 Subject: [PATCH 2/2] fix: scope results and query runs to a database --- CHANGELOG.md | 9 +++ README.md | 33 +++++--- examples/quickstart.rs | 82 +++++++++++++++----- src/arrow.rs | 43 ++++++++++- src/client.rs | 49 +++++++++--- src/query.rs | 144 +++++++++++++++++++++++++++++++++-- src/resources.rs | 33 +++++--- tests/query_async_polling.rs | 10 +-- tests/results_arrow.rs | 12 +-- 9 files changed, 345 insertions(+), 70 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc10e6d..221f04f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index b3dcc3c..c39b239 100644 --- a/README.md +++ b/README.md @@ -64,16 +64,22 @@ async fn main() -> Result<(), Box> { // .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,8 +192,11 @@ async fn main() -> Result<(), Box> { .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. - 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 { @@ -192,7 +204,7 @@ async fn main() -> Result<(), Box> { } // 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> { } ``` -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 diff --git a/examples/quickstart.rs b/examples/quickstart.rs index d95c9d0..5ff73b6 100644 --- a/examples/quickstart.rs +++ b/examples/quickstart.rs @@ -92,18 +92,36 @@ async fn run() -> Result<(), Box> { // --- 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> { ); // 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> { // 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, Box> { + 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> { +async fn fetch_arrow( + client: &Client, + database_id: &str, + result_id: &str, +) -> Result<(), Box> { 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 Result<(), Box> { +async fn one_shot_arrow( + client: &Client, + database_id: &str, +) -> Result<(), Box> { use std::time::Duration; println!("One-call query_to_arrow..."); @@ -194,6 +239,7 @@ async fn one_shot_arrow(client: &Client) -> Result<(), Box, limit: Option, ) -> Result { 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, limit: Option, ) -> Result { 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, limit: Option, ) -> Result<(Bytes, Option, Option), 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)])); diff --git a/src/client.rs b/src/client.rs index 314f08e..155ac06 100644 --- a/src/client.rs +++ b/src/client.rs @@ -680,32 +680,46 @@ impl Client { } } - /// List recent query runs. + /// List recent query runs for the database named by `database_id` (the + /// required `X-Database-Id` header — query runs are scoped to a database). pub async fn list_query_runs( &self, + database_id: &str, limit: Option, cursor: Option<&str>, ) -> Result> { - apis::query_runs_api::list_query_runs(&self.configuration, limit, cursor, None, None).await + apis::query_runs_api::list_query_runs( + &self.configuration, + database_id, + limit, + cursor, + None, + None, + ) + .await } - /// List persisted results. + /// List persisted results for the database named by `database_id` (the + /// required `X-Database-Id` header — results are scoped to a database). pub async fn list_results( &self, + database_id: &str, limit: Option, offset: Option, ) -> Result> { - apis::results_api::list_results(&self.configuration, limit, offset).await + apis::results_api::list_results(&self.configuration, database_id, limit, offset).await } - /// Fetch a persisted result by id (JSON form). For Arrow IPC decoding use + /// Fetch a persisted result by id (JSON form), scoped to `database_id` (the + /// required `X-Database-Id` header). For Arrow IPC decoding use /// [`Client::get_result_arrow`] (requires the `arrow` feature). pub async fn get_result( &self, id: &str, + database_id: &str, ) -> Result> { - apis::results_api::get_result(&self.configuration, id, None, None, None).await + apis::results_api::get_result(&self.configuration, id, database_id, None, None, None).await } /// List workspaces visible to the authenticated principal. @@ -726,10 +740,11 @@ impl Client { pub async fn get_result_arrow( &self, id: &str, + database_id: &str, offset: Option, limit: Option, ) -> Result { - crate::arrow::get_result_arrow(&self.configuration, id, offset, limit).await + crate::arrow::get_result_arrow(&self.configuration, id, database_id, offset, limit).await } /// Fetch a result as Arrow IPC and lazily iterate over its record batches. @@ -739,10 +754,11 @@ impl Client { pub async fn stream_result_arrow( &self, id: &str, + database_id: &str, offset: Option, limit: Option, ) -> Result { - crate::arrow::stream_result_arrow(&self.configuration, id, offset, limit).await + crate::arrow::stream_result_arrow(&self.configuration, id, database_id, offset, limit).await } // --- Resource handles ----------------------------------------------------- @@ -845,15 +861,19 @@ impl Client { /// A `failed` status returns [`AwaitResultError::Failed`]; exceeding /// `poll.timeout` returns [`AwaitResultError::Timeout`]. Use /// `PollConfig::default()` for sensible defaults (120s timeout, 1s interval). + /// + /// `database_id` scopes the result lookup to a database (the required + /// `X-Database-Id` header) — pass the same database the query ran in. pub async fn await_result( &self, result_id: &str, + database_id: &str, poll: PollConfig, ) -> Result { let deadline = std::time::Instant::now() + poll.timeout; loop { let result = self - .get_result(result_id) + .get_result(result_id, database_id) .await .map_err(AwaitResultError::Api)?; match crate::status::ResultStatus::parse(&result.status) { @@ -893,6 +913,7 @@ impl Client { pub async fn query_to_arrow( &self, request: models::QueryRequest, + database_id: &str, poll: PollConfig, offset: Option, limit: Option, @@ -900,8 +921,9 @@ impl Client { // Use the raw generated query (not the enhanced `Client::query`): the // Arrow path polls and fetches the persisted result itself, so it must // not also trigger truncation auto-follow, which would materialize the - // full result as JSON before we re-fetch it as Arrow. - let submitted = apis::query_api::query(&self.configuration, request, None) + // full result as JSON before we re-fetch it as Arrow. The query and the + // result fetch share `database_id` (the required `X-Database-Id` scope). + let submitted = apis::query_api::query(&self.configuration, request, Some(database_id)) .await .map_err(QueryToArrowError::Query)?; let result_id = @@ -919,7 +941,10 @@ impl Client { // second time as Arrow. let deadline = std::time::Instant::now() + poll.timeout; loop { - match self.get_result_arrow(&result_id, offset, limit).await { + match self + .get_result_arrow(&result_id, database_id, offset, limit) + .await + { Ok(result) => return Ok(result), Err(crate::arrow::ArrowError::NotReady { status, diff --git a/src/query.rs b/src/query.rs index e88d9f9..a5a7880 100644 --- a/src/query.rs +++ b/src/query.rs @@ -493,7 +493,16 @@ pub(crate) async fn execute_query( if !qc.auto_follow || !resp.truncated { Ok(resp) } else { - materialize_full(config, resp, qc).await + // Auto-follow re-fetches the persisted result and the query run, + // both of which are scoped to a database via the required + // `X-Database-Id` header. The spec requires exactly one database + // source on `/v1/query`, so the effective database is the header + // value when present, otherwise the request-body `database_id`. + let effective_db = x_database_id + .map(str::to_owned) + .or_else(|| request.database_id.clone().flatten()) + .unwrap_or_default(); + materialize_full(config, resp, &effective_db, qc).await } } } @@ -653,13 +662,17 @@ fn apply_apikey_headers( pub(crate) async fn wait_for_result( config: &Configuration, result_id: &str, + x_database_id: &str, poll: &PollPolicy, ) -> Result { let start = Instant::now(); let mut delay = poll.base_backoff; loop { // offset=None, limit=0 (status only), default format. - let result = match results_api::get_result(config, result_id, None, Some(0), None).await { + let result = + match results_api::get_result(config, result_id, x_database_id, None, Some(0), None) + .await + { Ok(result) => result, // A failed result is delivered as HTTP 409: the generated client // raises on any non-2xx rather than returning status="failed". The @@ -709,6 +722,7 @@ pub(crate) async fn wait_for_result( async fn materialize_full( config: &Configuration, mut preview: QueryResponse, + x_database_id: &str, qc: &QueryConfig, ) -> Result { let result_id = match preview.result_id.clone().flatten() { @@ -721,9 +735,9 @@ async fn materialize_full( } }; - wait_for_result(config, &result_id, &qc.poll).await?; + wait_for_result(config, &result_id, x_database_id, &qc.poll).await?; - let total = authoritative_total(config, &preview).await; + let total = authoritative_total(config, &preview, x_database_id).await; // Auto-follow does extra round-trips (poll + paginate) and materializes the // full result; log it so the hidden work behind one query() call is // observable without being noisy (info, not a warning). @@ -749,7 +763,7 @@ async fn materialize_full( } } - let rows = fetch_all_rows(config, &result_id, total, qc).await?; + let rows = fetch_all_rows(config, &result_id, x_database_id, total, qc).await?; // Replace the bounded preview with the full row set. `truncated` / // `total_row_count` stay as the server reported them so the caller can still @@ -766,11 +780,15 @@ async fn materialize_full( /// The grand total row count. `total_row_count` is null while a truncated result /// is still persisting, so fall back to the query-run record, which carries the /// authoritative count once the run succeeds; else unknown. -async fn authoritative_total(config: &Configuration, preview: &QueryResponse) -> Option { +async fn authoritative_total( + config: &Configuration, + preview: &QueryResponse, + x_database_id: &str, +) -> Option { if let Some(t) = preview.total_row_count.flatten() { return Some(t); } - match query_runs_api::get_query_run(config, &preview.query_run_id).await { + match query_runs_api::get_query_run(config, &preview.query_run_id, x_database_id).await { Ok(run) => run.row_count.flatten(), Err(_) => None, } @@ -786,6 +804,7 @@ async fn authoritative_total(config: &Configuration, preview: &QueryResponse) -> async fn fetch_all_rows( config: &Configuration, result_id: &str, + x_database_id: &str, total: Option, qc: &QueryConfig, ) -> Result>, QueryError> { @@ -797,6 +816,7 @@ async fn fetch_all_rows( let page = results_api::get_result( config, result_id, + x_database_id, Some(checked_offset(offset, result_id)?), Some(page_size), Some(ResultsFormatQuery::Json), @@ -922,7 +942,7 @@ mod tests { #[cfg(unix)] use crate::test_support::reset_then_ok_server; use serde_json::json; - use wiremock::matchers::{method, path, query_param}; + use wiremock::matchers::{header, method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; /// Build a client pointed at `base_url` with a static bearer (no JWT mint) @@ -1205,6 +1225,114 @@ mod tests { assert_eq!(resp.total_row_count.flatten(), Some(3)); } + /// Regression (#84 / database-scoped results): a truncated result followed + /// via `query_in` must carry the `X-Database-Id` header on every follow-up + /// request — the readiness poll, the query-run total lookup, and each page + /// fetch. The GET mocks match *only* when the header is present, so a missing + /// header 404s the request and fails the query. + #[tokio::test] + async fn auto_follow_forwards_database_id_header() { + let server = MockServer::start().await; + // Truncated with a null total, so the follow-up also hits /v1/query-runs. + Mock::given(method("POST")) + .and(path("/v1/query")) + .and(header("X-Database-Id", "db_x")) + .respond_with(ResponseTemplate::new(200).set_body_json(preview_json( + true, + Some("rslt1"), + None, + ))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/results/rslt1")) + .and(query_param("limit", "0")) + .and(header("X-Database-Id", "db_x")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "result_id": "rslt1", "status": "ready" + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/query-runs/qrun1")) + .and(header("X-Database-Id", "db_x")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "created_at": "2026-06-15T00:00:00Z", + "id": "qrun1", + "snapshot_id": "snap1", + "sql_hash": "h", + "sql_text": "SELECT 1 AS x", + "status": "succeeded", + "row_count": 1 + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/results/rslt1")) + .and(query_param("offset", "0")) + .and(query_param("limit", "2")) + .and(header("X-Database-Id", "db_x")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "result_id": "rslt1", "status": "ready", "rows": [[1]] + }))) + .mount(&server) + .await; + + let client = test_client(&server.uri(), fast_config()); + let resp = client + .query_in(req(), "db_x") + .await + .expect("auto-follow should forward X-Database-Id and succeed"); + assert_eq!(resp.rows.len(), 1); + assert_eq!(resp.total_row_count.flatten(), Some(1)); + } + + /// When no header is passed, the effective database for auto-follow falls + /// back to the request-body `database_id`: `query()` sends that as + /// `X-Database-Id` on the follow-up fetches. The follow-up GET mocks require + /// the header; the POST does not (the body carries the scope there). + #[tokio::test] + async fn auto_follow_uses_body_database_id_when_no_header() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/query")) + .respond_with(ResponseTemplate::new(200).set_body_json(preview_json( + true, + Some("rslt1"), + Some(1), + ))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/results/rslt1")) + .and(query_param("limit", "0")) + .and(header("X-Database-Id", "db_body")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "result_id": "rslt1", "status": "ready" + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/results/rslt1")) + .and(query_param("offset", "0")) + .and(query_param("limit", "2")) + .and(header("X-Database-Id", "db_body")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "result_id": "rslt1", "status": "ready", "rows": [[1]] + }))) + .mount(&server) + .await; + + let client = test_client(&server.uri(), fast_config()); + let mut request = req(); + request.database_id = Some(Some("db_body".to_owned())); + let resp = client + .query(request) + .await + .expect("auto-follow should fall back to the body database_id"); + assert_eq!(resp.rows.len(), 1); + } + #[tokio::test] async fn auto_follow_falls_back_to_query_run_total() { let server = MockServer::start().await; diff --git a/src/resources.rs b/src/resources.rs index a1f5011..084b3ad 100644 --- a/src/resources.rs +++ b/src/resources.rs @@ -486,25 +486,36 @@ impl<'a> QueryRunsApi<'a> { Self { config } } - /// Fetch a query run by id. + /// Fetch a query run by id. `x_database_id` scopes the lookup to a database + /// (the required `X-Database-Id` header). pub async fn get( &self, id: &str, + x_database_id: &str, ) -> Result> { - apis::query_runs_api::get_query_run(self.config, id).await + apis::query_runs_api::get_query_run(self.config, id, x_database_id).await } - /// List recent query runs. + /// List recent query runs for the database named by `x_database_id` (the + /// required `X-Database-Id` header). pub async fn list( &self, + x_database_id: &str, limit: Option, cursor: Option<&str>, status: Option<&str>, saved_query_id: Option<&str>, ) -> Result> { - apis::query_runs_api::list_query_runs(self.config, limit, cursor, status, saved_query_id) - .await + apis::query_runs_api::list_query_runs( + self.config, + x_database_id, + limit, + cursor, + status, + saved_query_id, + ) + .await } } @@ -521,24 +532,28 @@ impl<'a> ResultsApi<'a> { Self { config } } - /// Fetch a persisted result by id (JSON form). + /// Fetch a persisted result by id (JSON form). `x_database_id` scopes the + /// lookup to a database (the required `X-Database-Id` header). pub async fn get( &self, id: &str, + x_database_id: &str, offset: Option, limit: Option, format: Option, ) -> Result> { - apis::results_api::get_result(self.config, id, offset, limit, format).await + apis::results_api::get_result(self.config, id, x_database_id, offset, limit, format).await } - /// List persisted results. + /// List persisted results for the database named by `x_database_id` (the + /// required `X-Database-Id` header). pub async fn list( &self, + x_database_id: &str, limit: Option, offset: Option, ) -> Result> { - apis::results_api::list_results(self.config, limit, offset).await + apis::results_api::list_results(self.config, x_database_id, limit, offset).await } } diff --git a/tests/query_async_polling.rs b/tests/query_async_polling.rs index 56f70e0..7f99e5d 100644 --- a/tests/query_async_polling.rs +++ b/tests/query_async_polling.rs @@ -39,7 +39,7 @@ async fn query_async_polling() { let mut request = models::QueryRequest::new("SELECT 1 AS x".to_string()); request.r#async = Some(true); request.async_after_ms = Some(Some(1000)); - request.database_id = Some(Some(database_id)); + request.database_id = Some(Some(database_id.clone())); let outcome = client .submit_query(request, None) .await @@ -55,7 +55,7 @@ async fn query_async_polling() { let deadline = Instant::now() + POLL_TIMEOUT; let mut run: Option = None; while Instant::now() < deadline { - let current = query_runs_api::get_query_run(config, &query_run_id) + let current = query_runs_api::get_query_run(config, &query_run_id, &database_id) .await .expect("get_query_run should succeed"); let terminal = is_terminal(¤t.status); @@ -80,7 +80,7 @@ async fn query_async_polling() { ); assert_eq!(run.row_count, Some(Some(1))); - let runs_listing = query_runs_api::list_query_runs(config, Some(50), None, None, None) + let runs_listing = query_runs_api::list_query_runs(config, &database_id, Some(50), None, None, None) .await .expect("list_query_runs should succeed"); assert!( @@ -90,7 +90,7 @@ async fn query_async_polling() { if let Some(Some(result_id)) = run.result_id { let result = client - .get_result(&result_id) + .get_result(&result_id, &database_id) .await .expect("get_result should succeed"); assert_eq!(result.result_id, result_id); @@ -106,7 +106,7 @@ async fn query_async_polling() { // ResultInfo (list_results) exposes the id as `id`, not `result_id`. let results_listing = client - .list_results(Some(50), None) + .list_results(&database_id, Some(50), None) .await .expect("list_results should succeed"); assert!( diff --git a/tests/results_arrow.rs b/tests/results_arrow.rs index 3e6a1aa..04d5559 100644 --- a/tests/results_arrow.rs +++ b/tests/results_arrow.rs @@ -80,7 +80,7 @@ async fn results_arrow() { ); request.r#async = Some(true); request.async_after_ms = Some(Some(1000)); - request.database_id = Some(Some(database_id)); + request.database_id = Some(Some(database_id.clone())); // `submit_query` recovers the run id whether the query ran inline (HTTP 200) // or went async (HTTP 202); the enhanced `client.query` reports a 202 as // `QueryError::Async`, so the async submission path uses `submit_query`. @@ -99,7 +99,7 @@ async fn results_arrow() { let deadline = Instant::now() + POLL_TIMEOUT; let mut run: Option = None; while Instant::now() < deadline { - let current = query_runs_api::get_query_run(config, &query_run_id) + let current = query_runs_api::get_query_run(config, &query_run_id, &database_id) .await .expect("get_query_run should succeed"); let terminal = is_terminal(¤t.status); @@ -126,7 +126,7 @@ async fn results_arrow() { let mut ready = false; while Instant::now() < deadline { let result = client - .get_result(&result_id) + .get_result(&result_id, &database_id) .await .expect("get_result should succeed"); if result.status == "ready" { @@ -139,7 +139,7 @@ async fn results_arrow() { // Buffered: full set of RecordBatches. let buffered = client - .get_result_arrow(&result_id, None, None) + .get_result_arrow(&result_id, &database_id, None, None) .await .expect("get_result_arrow should succeed"); assert_eq!(total_rows(&buffered.batches), 2, "expected 2 rows"); @@ -168,7 +168,7 @@ async fn results_arrow() { // Streaming: same data via the per-batch iterator. let stream = client - .stream_result_arrow(&result_id, None, None) + .stream_result_arrow(&result_id, &database_id, None, None) .await .expect("stream_result_arrow should succeed"); let streamed: Vec = stream @@ -184,7 +184,7 @@ async fn results_arrow() { // Pagination forwards correctly: offset=1, limit=1 -> just the second row. let sliced = client - .get_result_arrow(&result_id, Some(1), Some(1)) + .get_result_arrow(&result_id, &database_id, Some(1), Some(1)) .await .expect("get_result_arrow with offset/limit should succeed"); assert_eq!(total_rows(&sliced.batches), 1);