From 3ff43dad373ed56977831fbd7573dc50e61f2c88 Mon Sep 17 00:00:00 2001 From: sdairs Date: Wed, 10 Jun 2026 21:46:40 +0100 Subject: [PATCH 1/5] Support OAuth bearer auth for cloud service query MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When authenticated via OAuth (cloud auth login), `cloud service query` now sends the user's bearer token directly to the Query API endpoint instead of looking up or auto-provisioning a per-service Query API key — provisioning requires write access an OAuth token doesn't have. The API-key flow is unchanged. `--no-auto-enable` is a documented no-op under OAuth, and a 404 from the query host maps to a clear "endpoint not enabled" error pointing at `cloud service query-endpoint create`. Library: `run_query` is refactored onto a shared request builder and a new `run_query_bearer` method uses the client's own bearer token (AuthMismatch on Basic-auth clients). The Query API host is now derived per environment from the client's base URL (`api.` → `queries.`) instead of hard-coding the prod host; `CLICKHOUSE_CLOUD_QUERY_HOST` still overrides, and `with_query_host` pins it programmatically (used by tests). The `auth-provider: custom` header marks a custom Query API key, so it is only sent on the Basic path. Tests: wiremock request-shape coverage for both run_query variants in the library, unit tests for query-host derivation, and CLI subprocess tests asserting the bearer header, the absence of provisioning calls/stored keys under OAuth, and the stored-key Basic path. Closes #247 (code ready; blocked on server-side bearer support at the query endpoint for live verification). Co-Authored-By: Claude Fable 5 --- README.md | 13 +- crates/clickhouse-cloud-api/src/client.rs | 182 +++++++++++++-- .../tests/run_query_test.rs | 134 +++++++++++ .../tests/spec_coverage_test.rs | 7 +- crates/clickhousectl/src/cloud/cli.rs | 12 +- crates/clickhousectl/src/cloud/commands.rs | 89 +++++--- .../tests/cli_request_shape_test.rs | 215 ++++++++++++++++++ 7 files changed, 590 insertions(+), 62 deletions(-) create mode 100644 crates/clickhouse-cloud-api/tests/run_query_test.rs diff --git a/README.md b/README.md index cd5337c..e26f8b3 100644 --- a/README.md +++ b/README.md @@ -285,7 +285,7 @@ clickhousectl cloud auth login This opens your browser for authentication via the OAuth device flow. Tokens are saved to `.clickhouse/tokens.json` (project-local). -> **Note:** OAuth tokens provide **read-only** access. You can list and inspect resources (organizations, services, backups, etc.) but cannot create, modify, or delete them. For write operations, use API key authentication. +> **Note:** OAuth tokens provide **read-only** access. You can list and inspect resources (organizations, services, backups, etc.) but cannot create, modify, or delete them. For write operations, use API key authentication. `cloud service query` works under OAuth too, running read-only SQL as your own identity — see [Query API auth modes](#query-api-auth-modes). ### API key/secret (required for write operations) @@ -495,17 +495,18 @@ clickhousectl cloud service delete --force | `--private-preview-terms-checked` | Accept private preview terms when required | | `--enable-core-dumps` | Enable or disable service core dump collection | -#### Query API auto-provisioning +#### Query API auth modes -The first time `cloud service query` runs against a service without a stored key, it provisions a Query API endpoint for that service and creates a dedicated API key bound to it. The key (`keyId`, `keySecret`, and `endpointId`) is stored in `.clickhouse/credentials.json` under `service_query_keys.`, alongside any user-level API key. Subsequent queries run SQL over HTTP using that key — no `clickhouse` binary and no service password required. The key is scoped to a single service, so it can read and write (SELECT, INSERT, DDL) against that service but cannot reach any other service in the org. Pass `--no-auto-enable` to fail instead of provisioning. +`cloud service query` is the canonical way to run SQL against a cloud service — over HTTP, with no `clickhouse` binary and no service password required. It works with both credential modes: + +- **API key auth** (read + write SQL): the first time `cloud service query` runs against a service without a stored key, it provisions a Query API endpoint for that service and creates a dedicated API key bound to it. The key (`keyId`, `keySecret`, and `endpointId`) is stored in `.clickhouse/credentials.json` under `service_query_keys.`, alongside any user-level API key. Subsequent queries use that key. It is scoped to a single service, so it can read and write (SELECT, INSERT, DDL) against that service but cannot reach any other service in the org. Pass `--no-auto-enable` to fail instead of provisioning. +- **OAuth** (`cloud auth login`, read-only SQL): the query runs as your own identity — the CLI sends your bearer token straight to the query endpoint. No Query API key is provisioned or stored (provisioning needs write access an OAuth token doesn't have), so the query endpoint must already be enabled on the service; an org admin can enable it with `cloud service query-endpoint create `. `--no-auto-enable` has no effect in this mode. Provisioning happens lazily (rather than at `service create` time) because the endpoint can only be bound once the service has finished provisioning, which can take several minutes — `service create` returns immediately instead of blocking on it. Per-service scoping is enforced at the query endpoint binding, which is created with role `sql_console_admin` (read + write inside the bound service only). The API key itself has no org-level roles, so the binding is the only thing that grants it any access. `cloud service delete` removes the stored key from `credentials.json`. -`cloud service query` is the canonical way to run SQL against a cloud service. - -Set `CLICKHOUSE_CLOUD_QUERY_HOST` to override the Query API host (defaults to `https://queries.clickhouse.cloud`). +The Query API host is derived from the API base URL per environment (`api.` → `queries.`, e.g. `https://queries.clickhouse.cloud` for production). Set `CLICKHOUSE_CLOUD_QUERY_HOST` to override it. ### Postgres (beta) diff --git a/crates/clickhouse-cloud-api/src/client.rs b/crates/clickhouse-cloud-api/src/client.rs index 597fb44..d4ce5cb 100644 --- a/crates/clickhouse-cloud-api/src/client.rs +++ b/crates/clickhouse-cloud-api/src/client.rs @@ -12,6 +12,18 @@ enum Auth { Bearer { token: String }, } +/// Credentials for a single Query API request. Basic carries a per-service +/// Query API key; Bearer carries the user's OAuth token. +enum QueryAuth<'a> { + Basic { + key_id: &'a str, + key_secret: &'a str, + }, + Bearer { + token: &'a str, + }, +} + /// ClickHouse Cloud API client. /// /// Supports both HTTP Basic Auth (API key/secret) and Bearer token (OAuth) authentication. @@ -20,6 +32,19 @@ pub struct Client { http: reqwest::Client, base_url: String, auth: Auth, + /// Explicit Query API host override; see [`Client::with_query_host`]. + query_host: Option, +} + +/// Derive the Query API host from a management API base URL by swapping the +/// `api.` host prefix for `queries.`, so each environment talks to its own +/// query host (e.g. `https://api.clickhouse-staging.com` → +/// `https://queries.clickhouse-staging.com`). Returns `None` when the base +/// URL isn't of that shape (e.g. a localhost test server). +fn derive_query_host(base_url: &str) -> Option { + let parsed = url::Url::parse(base_url).ok()?; + let rest = parsed.host_str()?.strip_prefix("api.")?; + Some(format!("{}://queries.{}", parsed.scheme(), rest)) } impl Client { @@ -41,6 +66,7 @@ impl Client { key_id: key_id.into(), key_secret: key_secret.into(), }, + query_host: None, } } @@ -55,6 +81,7 @@ impl Client { auth: Auth::Bearer { token: token.into(), }, + query_host: None, } } @@ -75,6 +102,7 @@ impl Client { key_id: key_id.into(), key_secret: key_secret.into(), }, + query_host: None, } } @@ -93,6 +121,7 @@ impl Client { auth: Auth::Bearer { token: token.into(), }, + query_host: None, } } @@ -112,6 +141,31 @@ impl Client { } } + /// Override the Query API host used by [`Client::run_query`] and + /// [`Client::run_query_bearer`]. + /// + /// When not set, the host is taken from the `CLICKHOUSE_CLOUD_QUERY_HOST` + /// env var if present, otherwise derived from the client's base URL + /// (`api.` → `queries.`), falling back to the production + /// host `https://queries.clickhouse.cloud`. + pub fn with_query_host(mut self, host: impl Into) -> Self { + self.query_host = Some(host.into().trim_end_matches('/').to_string()); + self + } + + /// Resolve the Query API host: explicit override, then env var, then + /// derivation from the base URL, then the production default. + fn resolved_query_host(&self) -> String { + if let Some(host) = &self.query_host { + return host.clone(); + } + if let Ok(host) = std::env::var("CLICKHOUSE_CLOUD_QUERY_HOST") { + return host; + } + derive_query_host(&self.base_url) + .unwrap_or_else(|| "https://queries.clickhouse.cloud".to_string()) + } + fn request(&self, method: reqwest::Method, path: &str) -> reqwest::RequestBuilder { let builder = self .http @@ -124,12 +178,12 @@ impl Client { /// Run a SQL statement against a service's Query API endpoint. /// - /// Hits `queries.clickhouse.cloud` (override via the - /// `CLICKHOUSE_CLOUD_QUERY_HOST` env var) using Basic auth with the - /// provided `key_id`/`key_secret` — a per-service key bound to a - /// query endpoint with role `sql_console_read_only` (or - /// `sql_console_admin`). This bypasses the client's primary auth - /// because Query API keys are scoped to a single service. + /// Hits the environment's query host (see [`Client::with_query_host`] + /// for resolution order) using Basic auth with the provided + /// `key_id`/`key_secret` — a per-service key bound to a query endpoint + /// with role `sql_console_read_only` (or `sql_console_admin`). This + /// bypasses the client's primary auth because Query API keys are scoped + /// to a single service. /// /// Returns the streaming response so the caller can forward it to /// stdout or buffer it into memory. @@ -141,6 +195,58 @@ impl Client { sql: &str, database: Option<&str>, format: &str, + ) -> Result { + self.run_query_with( + QueryAuth::Basic { key_id, key_secret }, + service_id, + sql, + database, + format, + ) + .await + } + + /// Run a SQL statement against a service's Query API endpoint using the + /// client's own OAuth Bearer token. + /// + /// Unlike [`Client::run_query`], no per-service Query API key is needed: + /// the query endpoint authenticates the user's identity directly, with + /// read-only SQL access. The query endpoint must already be enabled on + /// the service. + /// + /// Returns an error if the client is using Basic auth. + pub async fn run_query_bearer( + &self, + service_id: &str, + sql: &str, + database: Option<&str>, + format: &str, + ) -> Result { + let token = match &self.auth { + Auth::Bearer { token } => token, + Auth::Basic { .. } => { + return Err(Error::AuthMismatch( + "run_query_bearer called on a Basic-auth client".into(), + )); + } + }; + self.run_query_with( + QueryAuth::Bearer { token }, + service_id, + sql, + database, + format, + ) + .await + } + + async fn run_query_with( + &self, + auth: QueryAuth<'_>, + service_id: &str, + sql: &str, + database: Option<&str>, + format: &str, ) -> Result { #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] @@ -151,11 +257,9 @@ impl Client { database: Option<&'a str>, } - let host = std::env::var("CLICKHOUSE_CLOUD_QUERY_HOST") - .unwrap_or_else(|_| "https://queries.clickhouse.cloud".to_string()); let url = format!( "{}/service/{}/run", - host.trim_end_matches('/'), + self.resolved_query_host().trim_end_matches('/'), service_id, ); @@ -165,17 +269,23 @@ impl Client { database, }; - let response = self + let request = self .http .post(url) .query(&[("format", format)]) - .basic_auth(key_id, Some(key_secret)) .header("content-type", "text/plain;charset=UTF-8") - .header("x-service-type", "clickhouse") - .header("auth-provider", "custom") - .json(&body) - .send() - .await?; + .header("x-service-type", "clickhouse"); + // `auth-provider: custom` tells the query host the credentials are a + // custom (user-provisioned) Query API key. Bearer tokens carry their + // own provider information, so the header is omitted for them. + let request = match auth { + QueryAuth::Basic { key_id, key_secret } => request + .basic_auth(key_id, Some(key_secret)) + .header("auth-provider", "custom"), + QueryAuth::Bearer { token } => request.bearer_auth(token), + }; + + let response = request.json(&body).send().await?; let status = response.status(); if !status.is_success() { @@ -2702,4 +2812,44 @@ impl Client { Ok(serde_json::from_str(&body_text)?) } +} + +#[cfg(test)] +mod tests { + use super::derive_query_host; + + #[test] + fn derive_query_host_prod() { + assert_eq!( + derive_query_host("https://api.clickhouse.cloud").as_deref(), + Some("https://queries.clickhouse.cloud") + ); + } + + #[test] + fn derive_query_host_staging() { + assert_eq!( + derive_query_host("https://api.clickhouse-staging.com").as_deref(), + Some("https://queries.clickhouse-staging.com") + ); + } + + #[test] + fn derive_query_host_dev() { + assert_eq!( + derive_query_host("https://api.clickhouse-dev.com").as_deref(), + Some("https://queries.clickhouse-dev.com") + ); + } + + #[test] + fn derive_query_host_non_api_host_is_none() { + assert_eq!(derive_query_host("http://127.0.0.1:8123"), None); + assert_eq!(derive_query_host("https://example.com"), None); + } + + #[test] + fn derive_query_host_invalid_url_is_none() { + assert_eq!(derive_query_host("not a url"), None); + } } \ No newline at end of file diff --git a/crates/clickhouse-cloud-api/tests/run_query_test.rs b/crates/clickhouse-cloud-api/tests/run_query_test.rs new file mode 100644 index 0000000..d0cf49e --- /dev/null +++ b/crates/clickhouse-cloud-api/tests/run_query_test.rs @@ -0,0 +1,134 @@ +//! Request-shape tests for the Query API methods (`run_query` and +//! `run_query_bearer`) against a local wiremock server. +//! +//! These assert the auth header, request body shape, and headers each +//! variant puts on the wire, without touching any cloud infrastructure — +//! the real Query API is exercised by the cloud integration tests. The +//! query host is pinned with `with_query_host` so the tests are independent +//! of the `CLICKHOUSE_CLOUD_QUERY_HOST` env var and host derivation. + +use base64::Engine as _; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use clickhouse_cloud_api::{Client, Error}; + +async fn start_mock_query_host(status: u16, body: &str) -> MockServer { + let mock = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/service/svc-1/run")) + .respond_with(ResponseTemplate::new(status).set_body_string(body)) + .mount(&mock) + .await; + mock +} + +#[tokio::test] +async fn run_query_sends_basic_auth_with_query_key() { + let mock = start_mock_query_host(200, "1\n").await; + let client = + Client::with_base_url(mock.uri(), "org-key", "org-secret").with_query_host(mock.uri()); + + let response = client + .run_query( + "svc-1", + "query-key", + "query-secret", + "SELECT 1", + None, + "TabSeparated", + ) + .await + .expect("run_query failed"); + assert_eq!(response.status(), 200); + + let requests = mock.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + let request = &requests[0]; + + // Basic auth must use the per-service Query API key, not the client's + // primary (org-level) credentials. + let auth = request.headers.get("authorization").unwrap().to_str().unwrap(); + let expected = format!( + "Basic {}", + base64::engine::general_purpose::STANDARD.encode("query-key:query-secret") + ); + assert_eq!(auth, expected); + + assert_eq!(request.headers.get("auth-provider").unwrap(), "custom"); + assert_eq!(request.headers.get("x-service-type").unwrap(), "clickhouse"); + + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["sql"], "SELECT 1"); + assert!(body["runId"].as_str().is_some(), "runId missing: {body}"); + assert!( + body.get("database").is_none(), + "database leaked into body when not set: {body}" + ); + + let format = request + .url + .query_pairs() + .find(|(k, _)| k == "format") + .map(|(_, v)| v.to_string()); + assert_eq!(format.as_deref(), Some("TabSeparated")); +} + +#[tokio::test] +async fn run_query_bearer_sends_bearer_token() { + let mock = start_mock_query_host(200, "1\n").await; + let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); + + let response = client + .run_query_bearer("svc-1", "SELECT 1", Some("mydb"), "JSONEachRow") + .await + .expect("run_query_bearer failed"); + assert_eq!(response.status(), 200); + + let requests = mock.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + let request = &requests[0]; + + let auth = request.headers.get("authorization").unwrap().to_str().unwrap(); + assert_eq!(auth, "Bearer oauth-token"); + + // `auth-provider: custom` marks a custom Query API key; it must not be + // sent alongside a bearer token. + assert!(request.headers.get("auth-provider").is_none()); + assert_eq!(request.headers.get("x-service-type").unwrap(), "clickhouse"); + + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["sql"], "SELECT 1"); + assert_eq!(body["database"], "mydb"); +} + +#[tokio::test] +async fn run_query_bearer_on_basic_auth_client_is_auth_mismatch() { + let client = Client::with_base_url("https://api.clickhouse.cloud", "k", "s"); + let err = client + .run_query_bearer("svc-1", "SELECT 1", None, "CSV") + .await + .expect_err("expected AuthMismatch"); + assert!( + matches!(err, Error::AuthMismatch(_)), + "expected AuthMismatch, got: {err:?}" + ); +} + +#[tokio::test] +async fn run_query_non_success_status_maps_to_api_error() { + let mock = start_mock_query_host(404, "query endpoint not found").await; + let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); + + let err = client + .run_query_bearer("svc-1", "SELECT 1", None, "CSV") + .await + .expect_err("expected Api error"); + match err { + Error::Api { status, message } => { + assert_eq!(status, 404); + assert_eq!(message, "query endpoint not found"); + } + other => panic!("expected Error::Api, got: {other:?}"), + } +} diff --git a/crates/clickhouse-cloud-api/tests/spec_coverage_test.rs b/crates/clickhouse-cloud-api/tests/spec_coverage_test.rs index 39fc6ee..52a3d51 100644 --- a/crates/clickhouse-cloud-api/tests/spec_coverage_test.rs +++ b/crates/clickhouse-cloud-api/tests/spec_coverage_test.rs @@ -73,9 +73,10 @@ fn spec_schema_type_names(spec: &Value) -> BTreeSet { } /// Public client methods that intentionally don't correspond to an OpenAPI -/// operation. The Query API endpoint that backs `run_query` is hosted at -/// `queries.clickhouse.cloud` and is not described by the control-plane spec. -const NON_OPENAPI_CLIENT_METHODS: &[&str] = &["run_query"]; +/// operation. The Query API endpoint that backs `run_query` / +/// `run_query_bearer` is hosted at `queries.` (e.g. +/// `queries.clickhouse.cloud`) and is not described by the control-plane spec. +const NON_OPENAPI_CLIENT_METHODS: &[&str] = &["run_query", "run_query_bearer"]; fn assert_client_operation_coverage(spec: &Value) { let spec_operations = spec_operation_ids(spec); diff --git a/crates/clickhousectl/src/cloud/cli.rs b/crates/clickhousectl/src/cloud/cli.rs index d04f02f..75e2362 100644 --- a/crates/clickhousectl/src/cloud/cli.rs +++ b/crates/clickhousectl/src/cloud/cli.rs @@ -773,9 +773,12 @@ CONTEXT FOR AGENTS: #[command(after_help = "\ CONTEXT FOR AGENTS: Runs SQL over HTTP — no local clickhouse binary or service password required. - Uses a per-service API key (read+write, scoped to this service via the - query endpoint binding) auto-provisioned on first use and stored in - .clickhouse/credentials.json. + With API key auth: uses a per-service API key (read+write, scoped to this + service via the query endpoint binding) auto-provisioned on first use and + stored in .clickhouse/credentials.json. + With OAuth (cloud auth login): sends your own bearer token — read-only SQL, + no key provisioning; the query endpoint must already be enabled by an org + admin (cloud service query-endpoint create). SQL precedence: --query > --queries-file > stdin. Default format: PrettyCompact on a TTY, TabSeparated when piped.")] Query { @@ -808,7 +811,8 @@ CONTEXT FOR AGENTS: org_id: Option, /// Fail instead of auto-provisioning the query endpoint + API key - /// when none is stored locally + /// when none is stored locally (API key auth only; with OAuth nothing + /// is ever provisioned, so this flag has no effect) #[arg(long)] no_auto_enable: bool, }, diff --git a/crates/clickhousectl/src/cloud/commands.rs b/crates/clickhousectl/src/cloud/commands.rs index 0dacd8d..6c3325d 100644 --- a/crates/clickhousectl/src/cloud/commands.rs +++ b/crates/clickhousectl/src/cloud/commands.rs @@ -2075,43 +2075,66 @@ pub async fn service_query( .await?; let service_id = service.id.to_string(); - let key = match credentials::get_service_query_key(&service_id) { - Some(k) => k, - None if opts.no_auto_enable => { - return Err(format!( - "no stored Query API key for service {service_id}; rerun without --no-auto-enable to auto-provision" - ) - .into()); - } - None => { - eprintln!( - "Provisioning Query API endpoint + key for service '{}'...", - service.name - ); - crate::cloud::service_query::ensure_service_query_setup( - client, - &org_id, + let format = opts.format.unwrap_or_else(default_query_format); + + let response = if client.is_bearer_auth() { + // OAuth: the query endpoint authenticates the user's bearer token + // directly — read-only SQL, no per-service Query API key, and no + // auto-provisioning (which would need write access the token doesn't + // have). `--no-auto-enable` is a no-op here since nothing is ever + // provisioned. The endpoint itself must already be enabled. + client + .api() + .run_query_bearer(&service_id, &sql, opts.database.as_deref(), &format) + .await + .map_err(|e| -> Box { + match e { + clickhouse_cloud_api::Error::Api { status: 404, .. } => format!( + "the query endpoint for service '{}' is not enabled. Ask an org admin to enable it:\n clickhousectl cloud service query-endpoint create {service_id}", + service.name + ) + .into(), + e => client.convert_error(e).into(), + } + })? + } else { + let key = match credentials::get_service_query_key(&service_id) { + Some(k) => k, + None if opts.no_auto_enable => { + return Err(format!( + "no stored Query API key for service {service_id}; rerun without --no-auto-enable to auto-provision" + ) + .into()); + } + None => { + eprintln!( + "Provisioning Query API endpoint + key for service '{}'...", + service.name + ); + crate::cloud::service_query::ensure_service_query_setup( + client, + &org_id, + &service_id, + &service.name, + ) + .await? + } + }; + + client + .api() + .run_query( &service_id, - &service.name, + &key.key_id, + &key.key_secret, + &sql, + opts.database.as_deref(), + &format, ) - .await? - } + .await + .map_err(|e| client.convert_error(e))? }; - let format = opts.format.unwrap_or_else(default_query_format); - let response = client - .api() - .run_query( - &service_id, - &key.key_id, - &key.key_secret, - &sql, - opts.database.as_deref(), - &format, - ) - .await - .map_err(|e| client.convert_error(e))?; - use futures_util::StreamExt; use std::io::Write as _; let mut stream = response.bytes_stream(); diff --git a/crates/clickhousectl/tests/cli_request_shape_test.rs b/crates/clickhousectl/tests/cli_request_shape_test.rs index 7a425fe..f371d88 100644 --- a/crates/clickhousectl/tests/cli_request_shape_test.rs +++ b/crates/clickhousectl/tests/cli_request_shape_test.rs @@ -1246,6 +1246,221 @@ async fn dotenv_creds_produce_basic_auth_request() { ); } +// ── Service query auth modes (issue #247) ────────────────────────────────── +// +// `cloud service query` has two auth paths: +// - API key auth: a per-service Query API key (stored locally, else +// auto-provisioned) is sent as Basic auth to the query host. +// - OAuth: the user's own bearer token is sent directly to the query host +// — no key lookup and, crucially, NO provisioning calls (key creation +// and endpoint upsert need write access an OAuth token doesn't have). +// Both tests run the binary against two mocks: one impersonating the +// control plane (service lookup), one impersonating the query host (wired +// up via CLICKHOUSE_CLOUD_QUERY_HOST, which overrides host derivation). + +const QUERY_TEST_SERVICE_ID: &str = "11111111-2222-3333-4444-555555555555"; + +async fn start_mock_control_plane_with_service() -> MockServer { + let mock = MockServer::start().await; + let stub_service = serde_json::json!({ + "result": { "id": QUERY_TEST_SERVICE_ID, "name": "demo" }, + "status": 200, + "requestId": "stub-service-get", + }); + Mock::given(method("GET")) + .and(path(format!( + "/v1/organizations/org-1/services/{QUERY_TEST_SERVICE_ID}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(stub_service)) + .mount(&mock) + .await; + mock +} + +async fn start_mock_query_host() -> MockServer { + let mock = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/service/{QUERY_TEST_SERVICE_ID}/run"))) + .respond_with(ResponseTemplate::new(200).set_body_string("1\n")) + .mount(&mock) + .await; + mock +} + +fn assert_success(output: &std::process::Output) { + assert!( + output.status.success(), + "clickhousectl exited {}\nstderr:\n{}\nstdout:\n{}", + output.status.code().unwrap_or(-1), + String::from_utf8_lossy(&output.stderr), + String::from_utf8_lossy(&output.stdout), + ); +} + +#[tokio::test] +async fn service_query_with_oauth_sends_bearer_and_never_provisions() { + let control = start_mock_control_plane_with_service().await; + let query_host = start_mock_query_host().await; + + // OAuth tokens are the lowest-precedence credential tier, so clear the + // API key env vars and run in a temp dir whose .clickhouse/tokens.json + // is the only credential source. + let dir = tempfile::tempdir().unwrap(); + let ch_dir = dir.path().join(".clickhouse"); + std::fs::create_dir_all(&ch_dir).unwrap(); + let tokens = serde_json::json!({ + "access_token": "test-bearer-token", + "refresh_token": "unused", + "expires_at": 4102444800u64, // 2100-01-01: never expires in tests + "api_url": format!("{}/v1", control.uri()), + }); + std::fs::write( + ch_dir.join("tokens.json"), + serde_json::to_vec(&tokens).unwrap(), + ) + .unwrap(); + + let url = control.uri(); + let output = Command::new(clickhousectl_binary()) + .args([ + "cloud", + "--url", + &url, + "service", + "query", + "--id", + QUERY_TEST_SERVICE_ID, + "--org-id", + "org-1", + "--query", + "SELECT 1", + ]) + .current_dir(dir.path()) + .env_remove("CLICKHOUSE_CLOUD_API_KEY") + .env_remove("CLICKHOUSE_CLOUD_API_SECRET") + .env("CLICKHOUSE_CLOUD_QUERY_HOST", query_host.uri()) + .output() + .expect("failed to spawn clickhousectl"); + assert_success(&output); + + // The query request must carry the OAuth bearer token, not Basic auth, + // and no `auth-provider: custom` marker (that header means "custom + // Query API key"). + let query_requests = query_host.received_requests().await.unwrap(); + assert_eq!(query_requests.len(), 1); + let run = &query_requests[0]; + let auth = run.headers.get("authorization").unwrap().to_str().unwrap(); + assert_eq!(auth, "Bearer test-bearer-token"); + assert!( + run.headers.get("auth-provider").is_none(), + "auth-provider header must not accompany a bearer token", + ); + + // No provisioning: key creation and query-endpoint upsert are both + // POSTs, so the control plane must see only GETs. + let control_requests = control.received_requests().await.unwrap(); + assert!( + control_requests + .iter() + .all(|r| r.method == wiremock::http::Method::GET), + "OAuth service query made non-GET control-plane calls: {:?}", + control_requests + .iter() + .map(|r| format!("{} {}", r.method, r.url.path())) + .collect::>(), + ); + + // And no Query API key may be written locally for the OAuth path. + assert!( + !ch_dir.join("credentials.json").exists(), + "OAuth service query wrote .clickhouse/credentials.json", + ); + + // The query result streams through to stdout untouched. + assert_eq!(String::from_utf8_lossy(&output.stdout), "1\n"); +} + +#[tokio::test] +async fn service_query_with_stored_key_sends_basic_auth_with_that_key() { + let control = start_mock_control_plane_with_service().await; + let query_host = start_mock_query_host().await; + + // A stored per-service Query API key short-circuits provisioning; the + // control-plane creds come from the env tier (the credentials file + // carries only service_query_keys, no api_key/api_secret). + let dir = tempfile::tempdir().unwrap(); + let ch_dir = dir.path().join(".clickhouse"); + std::fs::create_dir_all(&ch_dir).unwrap(); + let creds = serde_json::json!({ + "service_query_keys": { + QUERY_TEST_SERVICE_ID: { + "key_id": "stored-key-id", + "key_secret": "stored-key-secret", + "endpoint_id": "ep-1", + "service_name": "demo", + "created_at": "2026-05-11T12:00:00Z", + } + } + }); + std::fs::write( + ch_dir.join("credentials.json"), + serde_json::to_vec(&creds).unwrap(), + ) + .unwrap(); + + let url = control.uri(); + let output = Command::new(clickhousectl_binary()) + .args([ + "cloud", + "--url", + &url, + "service", + "query", + "--id", + QUERY_TEST_SERVICE_ID, + "--org-id", + "org-1", + "--query", + "SELECT 1", + ]) + .current_dir(dir.path()) + .env("CLICKHOUSE_CLOUD_API_KEY", "fake-key-for-tests") + .env("CLICKHOUSE_CLOUD_API_SECRET", "fake-secret-for-tests") + .env("CLICKHOUSE_CLOUD_QUERY_HOST", query_host.uri()) + .output() + .expect("failed to spawn clickhousectl"); + assert_success(&output); + + // The query request authenticates with the stored per-service key, not + // the org-level env creds, and keeps the custom-key marker header. + let query_requests = query_host.received_requests().await.unwrap(); + assert_eq!(query_requests.len(), 1); + let run = &query_requests[0]; + let auth = run.headers.get("authorization").unwrap().to_str().unwrap(); + let expected = format!( + "Basic {}", + base64::Engine::encode( + &base64::engine::general_purpose::STANDARD, + "stored-key-id:stored-key-secret", + ) + ); + assert_eq!(auth, expected); + assert_eq!(run.headers.get("auth-provider").unwrap(), "custom"); + + // Stored key present → no provisioning POSTs on the control plane. + let control_requests = control.received_requests().await.unwrap(); + assert!( + control_requests + .iter() + .all(|r| r.method == wiremock::http::Method::GET), + "stored-key service query made non-GET control-plane calls: {:?}", + control_requests + .iter() + .map(|r| format!("{} {}", r.method, r.url.path())) + .collect::>(), + ); +} + // Shell env vars must win over `.env` — if both are set, the request is // signed with the shell values, never the file values. From 964a8d1e794281ce104884f8bc8e19b341fcc173 Mon Sep 17 00:00:00 2001 From: sdairs Date: Thu, 11 Jun 2026 09:35:24 +0100 Subject: [PATCH 2/5] Fix staging/dev API hosts and query-host derivation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The staging and dev entries in KNOWN_CONFIGS pointed at hosts that don't exist (api.clickhouse-staging.com / api.clickhouse-dev.com are NXDOMAIN); the real management API hosts carry a control-plane label: api.control-plane.clickhouse-staging.com and api.control-plane.clickhouse-dev.com. OAuth login against staging/dev could never have matched an auth config. Query-host derivation learns the same shape: the query hosts do NOT carry the control-plane label (queries.clickhouse-staging.com exists, queries.control-plane.clickhouse-staging.com doesn't), so derivation now strips an optional control-plane. prefix after api. — verified against live DNS for prod, staging, and dev. Co-Authored-By: Claude Fable 5 --- README.md | 2 +- crates/clickhouse-cloud-api/src/client.rs | 26 ++++++++++++++++++----- crates/clickhousectl/src/cloud/auth.rs | 15 +++++++------ 3 files changed, 31 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index e26f8b3..a3aa2d0 100644 --- a/README.md +++ b/README.md @@ -506,7 +506,7 @@ Provisioning happens lazily (rather than at `service create` time) because the e Per-service scoping is enforced at the query endpoint binding, which is created with role `sql_console_admin` (read + write inside the bound service only). The API key itself has no org-level roles, so the binding is the only thing that grants it any access. `cloud service delete` removes the stored key from `credentials.json`. -The Query API host is derived from the API base URL per environment (`api.` → `queries.`, e.g. `https://queries.clickhouse.cloud` for production). Set `CLICKHOUSE_CLOUD_QUERY_HOST` to override it. +The Query API host is derived from the API base URL per environment (`api.[control-plane.]` → `queries.`, e.g. `https://queries.clickhouse.cloud` for production). Set `CLICKHOUSE_CLOUD_QUERY_HOST` to override it. ### Postgres (beta) diff --git a/crates/clickhouse-cloud-api/src/client.rs b/crates/clickhouse-cloud-api/src/client.rs index d4ce5cb..53de4ba 100644 --- a/crates/clickhouse-cloud-api/src/client.rs +++ b/crates/clickhouse-cloud-api/src/client.rs @@ -38,12 +38,20 @@ pub struct Client { /// Derive the Query API host from a management API base URL by swapping the /// `api.` host prefix for `queries.`, so each environment talks to its own -/// query host (e.g. `https://api.clickhouse-staging.com` → -/// `https://queries.clickhouse-staging.com`). Returns `None` when the base -/// URL isn't of that shape (e.g. a localhost test server). +/// query host. Staging and dev serve the management API under an extra +/// `control-plane.` label that the query host doesn't have, so it is +/// dropped too: +/// +/// - `https://api.clickhouse.cloud` → `https://queries.clickhouse.cloud` +/// - `https://api.control-plane.clickhouse-staging.com` → +/// `https://queries.clickhouse-staging.com` +/// +/// Returns `None` when the base URL isn't of that shape (e.g. a localhost +/// test server). fn derive_query_host(base_url: &str) -> Option { let parsed = url::Url::parse(base_url).ok()?; let rest = parsed.host_str()?.strip_prefix("api.")?; + let rest = rest.strip_prefix("control-plane.").unwrap_or(rest); Some(format!("{}://queries.{}", parsed.scheme(), rest)) } @@ -2829,7 +2837,7 @@ mod tests { #[test] fn derive_query_host_staging() { assert_eq!( - derive_query_host("https://api.clickhouse-staging.com").as_deref(), + derive_query_host("https://api.control-plane.clickhouse-staging.com").as_deref(), Some("https://queries.clickhouse-staging.com") ); } @@ -2837,11 +2845,19 @@ mod tests { #[test] fn derive_query_host_dev() { assert_eq!( - derive_query_host("https://api.clickhouse-dev.com").as_deref(), + derive_query_host("https://api.control-plane.clickhouse-dev.com").as_deref(), Some("https://queries.clickhouse-dev.com") ); } + #[test] + fn derive_query_host_plain_api_prefix_without_control_plane() { + assert_eq!( + derive_query_host("https://api.clickhouse-staging.com").as_deref(), + Some("https://queries.clickhouse-staging.com") + ); + } + #[test] fn derive_query_host_non_api_host_is_none() { assert_eq!(derive_query_host("http://127.0.0.1:8123"), None); diff --git a/crates/clickhousectl/src/cloud/auth.rs b/crates/clickhousectl/src/cloud/auth.rs index 5168c6d..91181f6 100644 --- a/crates/clickhousectl/src/cloud/auth.rs +++ b/crates/clickhousectl/src/cloud/auth.rs @@ -21,14 +21,14 @@ const KNOWN_CONFIGS: &[(&str, AuthConfig)] = &[ }, ), ( - "api.clickhouse-staging.com", + "api.control-plane.clickhouse-staging.com", AuthConfig { auth_url: "https://auth.control-plane.clickhouse-staging.com", client_id: "ZC8AupPshQt2UNO2hEDutnKitx4PhizY", }, ), ( - "api.clickhouse-dev.com", + "api.control-plane.clickhouse-dev.com", AuthConfig { auth_url: "https://auth.control-plane.clickhouse-dev.com", client_id: "bVVcrqNw1t5dya9WFzfnM7PSsAgmfzwY", @@ -417,14 +417,17 @@ mod tests { #[test] fn test_auth_config_lookup() { assert!(auth_config_for_url("https://api.clickhouse.cloud/v1").is_some()); - assert!(auth_config_for_url("https://api.clickhouse-staging.com/v1").is_some()); - assert!(auth_config_for_url("https://api.clickhouse-dev.com/v1").is_some()); + assert!( + auth_config_for_url("https://api.control-plane.clickhouse-staging.com/v1").is_some() + ); + assert!(auth_config_for_url("https://api.control-plane.clickhouse-dev.com/v1").is_some()); assert!(auth_config_for_url("https://api.unknown.com/v1").is_none()); // Verify distinct configs let prod = auth_config_for_url("https://api.clickhouse.cloud/v1").unwrap(); - let staging = auth_config_for_url("https://api.clickhouse-staging.com/v1").unwrap(); - let dev = auth_config_for_url("https://api.clickhouse-dev.com/v1").unwrap(); + let staging = + auth_config_for_url("https://api.control-plane.clickhouse-staging.com/v1").unwrap(); + let dev = auth_config_for_url("https://api.control-plane.clickhouse-dev.com/v1").unwrap(); assert_ne!(prod.client_id, staging.client_id); assert_ne!(prod.client_id, dev.client_id); assert_ne!(staging.client_id, dev.client_id); From d3993dfff243efb109aba9b4c85f7674648b2bb1 Mon Sep 17 00:00:00 2001 From: sdairs Date: Thu, 11 Jun 2026 21:56:12 +0100 Subject: [PATCH 3/5] OAuth service query needs no query endpoint; align docs with live behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live testing against staging (after PR 34114 deployed) confirmed the JWT-authenticated Query API path is SQL-console style: it authenticates the user's own identity and never consults the service's query-endpoint configuration (verified: queries succeed against a service whose query-endpoint GET returns NOT_FOUND). Drop the 404 → "ask an admin to run query-endpoint create" error mapping, which was based on a wrong premise, and fix the help text, README, and run_query_bearer docs that claimed the endpoint must be pre-enabled. Also soften "read-only SQL" to "permissions follow your console role": the backend forces readonly only for the mcp/librechat strategies, not clickhousectl. Co-Authored-By: Claude Fable 5 --- README.md | 4 ++-- crates/clickhouse-cloud-api/src/client.rs | 8 ++++---- crates/clickhousectl/src/cloud/cli.rs | 6 +++--- crates/clickhousectl/src/cloud/commands.rs | 22 +++++++--------------- 4 files changed, 16 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index a3aa2d0..2acce8a 100644 --- a/README.md +++ b/README.md @@ -285,7 +285,7 @@ clickhousectl cloud auth login This opens your browser for authentication via the OAuth device flow. Tokens are saved to `.clickhouse/tokens.json` (project-local). -> **Note:** OAuth tokens provide **read-only** access. You can list and inspect resources (organizations, services, backups, etc.) but cannot create, modify, or delete them. For write operations, use API key authentication. `cloud service query` works under OAuth too, running read-only SQL as your own identity — see [Query API auth modes](#query-api-auth-modes). +> **Note:** OAuth tokens provide **read-only** access. You can list and inspect resources (organizations, services, backups, etc.) but cannot create, modify, or delete them. For write operations, use API key authentication. `cloud service query` works under OAuth too, running SQL as your own identity with your console role's permissions — see [Query API auth modes](#query-api-auth-modes). ### API key/secret (required for write operations) @@ -500,7 +500,7 @@ clickhousectl cloud service delete --force `cloud service query` is the canonical way to run SQL against a cloud service — over HTTP, with no `clickhouse` binary and no service password required. It works with both credential modes: - **API key auth** (read + write SQL): the first time `cloud service query` runs against a service without a stored key, it provisions a Query API endpoint for that service and creates a dedicated API key bound to it. The key (`keyId`, `keySecret`, and `endpointId`) is stored in `.clickhouse/credentials.json` under `service_query_keys.`, alongside any user-level API key. Subsequent queries use that key. It is scoped to a single service, so it can read and write (SELECT, INSERT, DDL) against that service but cannot reach any other service in the org. Pass `--no-auto-enable` to fail instead of provisioning. -- **OAuth** (`cloud auth login`, read-only SQL): the query runs as your own identity — the CLI sends your bearer token straight to the query endpoint. No Query API key is provisioned or stored (provisioning needs write access an OAuth token doesn't have), so the query endpoint must already be enabled on the service; an org admin can enable it with `cloud service query-endpoint create `. `--no-auto-enable` has no effect in this mode. +- **OAuth** (`cloud auth login`): the query runs as your own identity, SQL-console style — the CLI sends your bearer token straight to the Query API, and your SQL permissions follow your ClickHouse Cloud console role. No Query API key is provisioned or stored, and no query endpoint needs to be configured on the service. `--no-auto-enable` has no effect in this mode. Provisioning happens lazily (rather than at `service create` time) because the endpoint can only be bound once the service has finished provisioning, which can take several minutes — `service create` returns immediately instead of blocking on it. diff --git a/crates/clickhouse-cloud-api/src/client.rs b/crates/clickhouse-cloud-api/src/client.rs index 53de4ba..ddd8568 100644 --- a/crates/clickhouse-cloud-api/src/client.rs +++ b/crates/clickhouse-cloud-api/src/client.rs @@ -217,10 +217,10 @@ impl Client { /// Run a SQL statement against a service's Query API endpoint using the /// client's own OAuth Bearer token. /// - /// Unlike [`Client::run_query`], no per-service Query API key is needed: - /// the query endpoint authenticates the user's identity directly, with - /// read-only SQL access. The query endpoint must already be enabled on - /// the service. + /// Unlike [`Client::run_query`], no per-service Query API key and no + /// query-endpoint configuration are needed: the Query API authenticates + /// the user's identity directly (SQL-console style), and SQL permissions + /// follow the user's console role. /// /// Returns an error if the client is using Basic auth. pub async fn run_query_bearer( diff --git a/crates/clickhousectl/src/cloud/cli.rs b/crates/clickhousectl/src/cloud/cli.rs index 75e2362..66ac803 100644 --- a/crates/clickhousectl/src/cloud/cli.rs +++ b/crates/clickhousectl/src/cloud/cli.rs @@ -776,9 +776,9 @@ CONTEXT FOR AGENTS: With API key auth: uses a per-service API key (read+write, scoped to this service via the query endpoint binding) auto-provisioned on first use and stored in .clickhouse/credentials.json. - With OAuth (cloud auth login): sends your own bearer token — read-only SQL, - no key provisioning; the query endpoint must already be enabled by an org - admin (cloud service query-endpoint create). + With OAuth (cloud auth login): sends your own bearer token — SQL runs as + your cloud user (permissions follow your console role); no key provisioning + and no query endpoint required on the service. SQL precedence: --query > --queries-file > stdin. Default format: PrettyCompact on a TTY, TabSeparated when piped.")] Query { diff --git a/crates/clickhousectl/src/cloud/commands.rs b/crates/clickhousectl/src/cloud/commands.rs index 6c3325d..489c3b2 100644 --- a/crates/clickhousectl/src/cloud/commands.rs +++ b/crates/clickhousectl/src/cloud/commands.rs @@ -2078,25 +2078,17 @@ pub async fn service_query( let format = opts.format.unwrap_or_else(default_query_format); let response = if client.is_bearer_auth() { - // OAuth: the query endpoint authenticates the user's bearer token - // directly — read-only SQL, no per-service Query API key, and no - // auto-provisioning (which would need write access the token doesn't - // have). `--no-auto-enable` is a no-op here since nothing is ever - // provisioned. The endpoint itself must already be enabled. + // OAuth: the Query API authenticates the user's bearer token + // directly, SQL-console style — the query runs as the user's own + // cloud identity, with no per-service Query API key and no + // query-endpoint configuration needed on the service. + // `--no-auto-enable` is a no-op here since nothing is ever + // provisioned. client .api() .run_query_bearer(&service_id, &sql, opts.database.as_deref(), &format) .await - .map_err(|e| -> Box { - match e { - clickhouse_cloud_api::Error::Api { status: 404, .. } => format!( - "the query endpoint for service '{}' is not enabled. Ask an org admin to enable it:\n clickhousectl cloud service query-endpoint create {service_id}", - service.name - ) - .into(), - e => client.convert_error(e).into(), - } - })? + .map_err(|e| client.convert_error(e))? } else { let key = match credentials::get_service_query_key(&service_id) { Some(k) => k, From 6372966207cc40be7a3dc7e9cae88704132161b8 Mon Sep 17 00:00:00 2001 From: sdairs Date: Thu, 11 Jun 2026 22:46:13 +0100 Subject: [PATCH 4/5] Address review: preserve port in derive_query_host, refresh stale test hosts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - derive_query_host now keeps a non-default port from the base URL (api.mycorp.example.com:8443 → queries.mycorp.example.com:8443), with unit coverage for both custom and default ports. - Tests still referencing the dead api.clickhouse-staging.com host now use api.control-plane.clickhouse-staging.com (token serialization, URL normalization, lib_base_url). The plain-api-prefix derivation test keeps the old shape deliberately. - assert_success moved to the shared helpers in cli_request_shape_test.rs and reused by the dotenv and shell-env-precedence tests. - Restore trailing newline at end of client.rs. Co-Authored-By: Claude Fable 5 --- crates/clickhouse-cloud-api/src/client.rs | 22 +++++++++++-- crates/clickhousectl/src/cloud/auth.rs | 11 ++++--- crates/clickhousectl/src/cloud/client.rs | 4 +-- .../tests/cli_request_shape_test.rs | 32 ++++++++----------- 4 files changed, 43 insertions(+), 26 deletions(-) diff --git a/crates/clickhouse-cloud-api/src/client.rs b/crates/clickhouse-cloud-api/src/client.rs index ddd8568..9d47c9d 100644 --- a/crates/clickhouse-cloud-api/src/client.rs +++ b/crates/clickhouse-cloud-api/src/client.rs @@ -52,7 +52,11 @@ fn derive_query_host(base_url: &str) -> Option { let parsed = url::Url::parse(base_url).ok()?; let rest = parsed.host_str()?.strip_prefix("api.")?; let rest = rest.strip_prefix("control-plane.").unwrap_or(rest); - Some(format!("{}://queries.{}", parsed.scheme(), rest)) + let port = parsed + .port() + .map(|p| format!(":{p}")) + .unwrap_or_default(); + Some(format!("{}://queries.{}{}", parsed.scheme(), rest, port)) } impl Client { @@ -2868,4 +2872,18 @@ mod tests { fn derive_query_host_invalid_url_is_none() { assert_eq!(derive_query_host("not a url"), None); } -} \ No newline at end of file + + #[test] + fn derive_query_host_preserves_non_default_port() { + assert_eq!( + derive_query_host("https://api.mycorp.example.com:8443").as_deref(), + Some("https://queries.mycorp.example.com:8443") + ); + // Default ports are normalized away by the URL parser and stay off + // the derived host. + assert_eq!( + derive_query_host("https://api.clickhouse.cloud:443").as_deref(), + Some("https://queries.clickhouse.cloud") + ); + } +} diff --git a/crates/clickhousectl/src/cloud/auth.rs b/crates/clickhousectl/src/cloud/auth.rs index 91181f6..49a536c 100644 --- a/crates/clickhousectl/src/cloud/auth.rs +++ b/crates/clickhousectl/src/cloud/auth.rs @@ -364,12 +364,15 @@ mod tests { access_token: "a".into(), refresh_token: "r".into(), expires_at: 1700000000, - api_url: "https://api.clickhouse-staging.com/v1".into(), + api_url: "https://api.control-plane.clickhouse-staging.com/v1".into(), }; let json = serde_json::to_string(&tokens).unwrap(); assert!(json.contains("clickhouse-staging.com")); let parsed: TokenStore = serde_json::from_str(&json).unwrap(); - assert_eq!(parsed.api_url, "https://api.clickhouse-staging.com/v1"); + assert_eq!( + parsed.api_url, + "https://api.control-plane.clickhouse-staging.com/v1" + ); } #[test] @@ -453,8 +456,8 @@ mod tests { "https://api.clickhouse.cloud/v1" ); assert_eq!( - normalize_api_url("https://api.clickhouse-staging.com/v1/"), - "https://api.clickhouse-staging.com/v1" + normalize_api_url("https://api.control-plane.clickhouse-staging.com/v1/"), + "https://api.control-plane.clickhouse-staging.com/v1" ); } } diff --git a/crates/clickhousectl/src/cloud/client.rs b/crates/clickhousectl/src/cloud/client.rs index 2e98704..0493f23 100644 --- a/crates/clickhousectl/src/cloud/client.rs +++ b/crates/clickhousectl/src/cloud/client.rs @@ -1141,8 +1141,8 @@ mod tests { #[test] fn lib_base_url_strips_v1_from_staging() { assert_eq!( - lib_base_url("https://api.clickhouse-staging.com/v1"), - "https://api.clickhouse-staging.com" + lib_base_url("https://api.control-plane.clickhouse-staging.com/v1"), + "https://api.control-plane.clickhouse-staging.com" ); } diff --git a/crates/clickhousectl/tests/cli_request_shape_test.rs b/crates/clickhousectl/tests/cli_request_shape_test.rs index f371d88..06fa29a 100644 --- a/crates/clickhousectl/tests/cli_request_shape_test.rs +++ b/crates/clickhousectl/tests/cli_request_shape_test.rs @@ -61,6 +61,18 @@ async fn start_mock_clickpipes_api() -> MockServer { mock } +/// Assert the binary exited zero, panicking with the captured stderr/stdout +/// so the failure cause is visible in the test output. +fn assert_success(output: &std::process::Output) { + assert!( + output.status.success(), + "clickhousectl exited {}\nstderr:\n{}\nstdout:\n{}", + output.status.code().unwrap_or(-1), + String::from_utf8_lossy(&output.stderr), + String::from_utf8_lossy(&output.stdout), + ); +} + /// Run the clickhousectl binary against the mock, returning the JSON body /// the binary POSTed. Panics with the captured stderr if the binary exits /// non-zero — a failure here is almost always a clap-parsing error, which @@ -1215,13 +1227,7 @@ async fn dotenv_creds_produce_basic_auth_request() { .output() .expect("failed to spawn clickhousectl"); - assert!( - output.status.success(), - "clickhousectl exited {}\nstderr:\n{}\nstdout:\n{}", - output.status.code().unwrap_or(-1), - String::from_utf8_lossy(&output.stderr), - String::from_utf8_lossy(&output.stdout), - ); + assert_success(&output); let requests = mock .received_requests() @@ -1287,16 +1293,6 @@ async fn start_mock_query_host() -> MockServer { mock } -fn assert_success(output: &std::process::Output) { - assert!( - output.status.success(), - "clickhousectl exited {}\nstderr:\n{}\nstdout:\n{}", - output.status.code().unwrap_or(-1), - String::from_utf8_lossy(&output.stderr), - String::from_utf8_lossy(&output.stdout), - ); -} - #[tokio::test] async fn service_query_with_oauth_sends_bearer_and_never_provisions() { let control = start_mock_control_plane_with_service().await; @@ -1497,7 +1493,7 @@ async fn shell_env_overrides_dotenv_creds_in_request() { .output() .expect("failed to spawn clickhousectl"); - assert!(output.status.success(), "binary failed: {}", String::from_utf8_lossy(&output.stderr)); + assert_success(&output); let requests = mock.received_requests().await.unwrap(); let auth = requests From ae8acad7c5b562a659e0084612631b679405f168 Mon Sep 17 00:00:00 2001 From: sdairs Date: Thu, 11 Jun 2026 23:26:59 +0100 Subject: [PATCH 5/5] Auto-wake idled services for OAuth service query MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The SQL-console-style query endpoint the OAuth path uses does not wake an idled service on its own: it answers 206 {"data":"Confirm wake service"} and expects the query to be resent with a `wake-service: true` header once the user confirms (this is what the SQL console does after prompting). The CLI previously streamed that 206 body straight to stdout as if it were the query result. The library now maps the query host's 206 service-state protocol to typed errors (ServiceIdle / ServiceStopped) and run_query{,_bearer} gain a wake_service flag that sends the wake confirmation header. The CLI retries once with the flag set after printing a notice to stderr — matching the API-key path, which the query host wakes automatically — and turns ServiceStopped into a hint to run `cloud service start` (stopped services are never woken by the Query API). Co-Authored-By: Claude Fable 5 --- README.md | 2 + crates/clickhouse-cloud-api/src/client.rs | 42 +++++ crates/clickhouse-cloud-api/src/error.rs | 12 ++ .../tests/integration_test.rs | 4 + .../tests/run_query_test.rs | 81 ++++++++- crates/clickhousectl/src/cloud/commands.rs | 61 +++++-- .../tests/cli_request_shape_test.rs | 155 +++++++++++++++++- 7 files changed, 343 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 2acce8a..2d3abcf 100644 --- a/README.md +++ b/README.md @@ -506,6 +506,8 @@ Provisioning happens lazily (rather than at `service create` time) because the e Per-service scoping is enforced at the query endpoint binding, which is created with role `sql_console_admin` (read + write inside the bound service only). The API key itself has no org-level roles, so the binding is the only thing that grants it any access. `cloud service delete` removes the stored key from `credentials.json`. +Querying an **idled** service wakes it automatically in both auth modes — under OAuth the Query API first asks for a wake confirmation, which the CLI sends after printing a notice to stderr (the first query may take a minute while the service wakes). A **stopped** service is never woken: the query fails with a hint to run `cloud service start`. + The Query API host is derived from the API base URL per environment (`api.[control-plane.]` → `queries.`, e.g. `https://queries.clickhouse.cloud` for production). Set `CLICKHOUSE_CLOUD_QUERY_HOST` to override it. ### Postgres (beta) diff --git a/crates/clickhouse-cloud-api/src/client.rs b/crates/clickhouse-cloud-api/src/client.rs index 9d47c9d..3e033e4 100644 --- a/crates/clickhouse-cloud-api/src/client.rs +++ b/crates/clickhouse-cloud-api/src/client.rs @@ -197,8 +197,12 @@ impl Client { /// bypasses the client's primary auth because Query API keys are scoped /// to a single service. /// + /// `wake_service` resends the wake confirmation the query host asks for + /// when the target service is idled — see [`Error::ServiceIdle`]. + /// /// Returns the streaming response so the caller can forward it to /// stdout or buffer it into memory. + #[allow(clippy::too_many_arguments)] pub async fn run_query( &self, service_id: &str, @@ -207,6 +211,7 @@ impl Client { sql: &str, database: Option<&str>, format: &str, + wake_service: bool, ) -> Result { self.run_query_with( QueryAuth::Basic { key_id, key_secret }, @@ -214,6 +219,7 @@ impl Client { sql, database, format, + wake_service, ) .await } @@ -226,6 +232,9 @@ impl Client { /// the user's identity directly (SQL-console style), and SQL permissions /// follow the user's console role. /// + /// `wake_service` resends the wake confirmation the query host asks for + /// when the target service is idled — see [`Error::ServiceIdle`]. + /// /// Returns an error if the client is using Basic auth. pub async fn run_query_bearer( &self, @@ -233,6 +242,7 @@ impl Client { sql: &str, database: Option<&str>, format: &str, + wake_service: bool, ) -> Result { let token = match &self.auth { Auth::Bearer { token } => token, @@ -248,6 +258,7 @@ impl Client { sql, database, format, + wake_service, ) .await } @@ -259,6 +270,7 @@ impl Client { sql: &str, database: Option<&str>, format: &str, + wake_service: bool, ) -> Result { #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] @@ -287,6 +299,14 @@ impl Client { .query(&[("format", format)]) .header("content-type", "text/plain;charset=UTF-8") .header("x-service-type", "clickhouse"); + // `wake-service: true` is the wake confirmation the query host asks + // for via a 206 `Confirm wake service` response (the SQL console + // sends it after prompting the user). + let request = if wake_service { + request.header("wake-service", "true") + } else { + request + }; // `auth-provider: custom` tells the query host the credentials are a // custom (user-provisioned) Query API key. Bearer tokens carry their // own provider information, so the header is omitted for them. @@ -300,6 +320,28 @@ impl Client { let response = request.json(&body).send().await?; let status = response.status(); + // 206 means the service can't take the query in its current state: + // `Confirm wake service` for an idled service (resend with the + // wake confirmation to wake it and run the query), `Service is + // stopped` for one that must be started explicitly. + if status.as_u16() == 206 { + let body_text = response.text().await.unwrap_or_default(); + #[derive(serde::Deserialize)] + struct StateBody { + data: Option, + } + let data = serde_json::from_str::(&body_text) + .ok() + .and_then(|b| b.data); + return Err(match data.as_deref() { + Some("Confirm wake service") => Error::ServiceIdle, + Some("Service is stopped") => Error::ServiceStopped, + _ => Error::Api { + status: 206, + message: body_text, + }, + }); + } if !status.is_success() { let body_text = response.text().await.unwrap_or_default(); return Err(Error::Api { diff --git a/crates/clickhouse-cloud-api/src/error.rs b/crates/clickhouse-cloud-api/src/error.rs index ae05908..6fba34c 100644 --- a/crates/clickhouse-cloud-api/src/error.rs +++ b/crates/clickhouse-cloud-api/src/error.rs @@ -18,4 +18,16 @@ pub enum Error { /// Operation requires a different auth mode than the client was configured with. #[error("auth mismatch: {0}")] AuthMismatch(String), + + /// The Query API reported the service is idled and asked for an explicit + /// wake confirmation (HTTP 206 `Confirm wake service`). Retry the query + /// with `wake_service` set to wake the service and run it. + #[error("service is idle; retry the query with wake_service to wake it")] + ServiceIdle, + + /// The Query API reported the service is stopped (HTTP 206 `Service is + /// stopped`). A stopped service is never woken by the Query API; it must + /// be started explicitly. + #[error("service is stopped; it must be started before it can be queried")] + ServiceStopped, } diff --git a/crates/clickhouse-cloud-api/tests/integration_test.rs b/crates/clickhouse-cloud-api/tests/integration_test.rs index 6cf4411..4323736 100644 --- a/crates/clickhouse-cloud-api/tests/integration_test.rs +++ b/crates/clickhouse-cloud-api/tests/integration_test.rs @@ -311,6 +311,7 @@ async fn cloud_service_crud_lifecycle() -> TestResult<()> { "SELECT 1", None, "TabSeparated", + false, ) .await { @@ -455,6 +456,7 @@ async fn cloud_service_crud_lifecycle() -> TestResult<()> { "SELECT 1", None, "TabSeparated", + false, ) .await { @@ -522,6 +524,7 @@ async fn cloud_service_crud_lifecycle() -> TestResult<()> { sql, None, "TabSeparated", + false, ) .await?; response @@ -657,6 +660,7 @@ async fn cloud_service_crud_lifecycle() -> TestResult<()> { "SELECT 1", None, "TabSeparated", + false, ) .await?; let body = response diff --git a/crates/clickhouse-cloud-api/tests/run_query_test.rs b/crates/clickhouse-cloud-api/tests/run_query_test.rs index d0cf49e..b40fe0a 100644 --- a/crates/clickhouse-cloud-api/tests/run_query_test.rs +++ b/crates/clickhouse-cloud-api/tests/run_query_test.rs @@ -37,6 +37,7 @@ async fn run_query_sends_basic_auth_with_query_key() { "SELECT 1", None, "TabSeparated", + false, ) .await .expect("run_query failed"); @@ -57,6 +58,10 @@ async fn run_query_sends_basic_auth_with_query_key() { assert_eq!(request.headers.get("auth-provider").unwrap(), "custom"); assert_eq!(request.headers.get("x-service-type").unwrap(), "clickhouse"); + assert!( + request.headers.get("wake-service").is_none(), + "wake-service header must not be sent unless wake_service is set", + ); let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); assert_eq!(body["sql"], "SELECT 1"); @@ -80,7 +85,7 @@ async fn run_query_bearer_sends_bearer_token() { let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); let response = client - .run_query_bearer("svc-1", "SELECT 1", Some("mydb"), "JSONEachRow") + .run_query_bearer("svc-1", "SELECT 1", Some("mydb"), "JSONEachRow", false) .await .expect("run_query_bearer failed"); assert_eq!(response.status(), 200); @@ -106,7 +111,7 @@ async fn run_query_bearer_sends_bearer_token() { async fn run_query_bearer_on_basic_auth_client_is_auth_mismatch() { let client = Client::with_base_url("https://api.clickhouse.cloud", "k", "s"); let err = client - .run_query_bearer("svc-1", "SELECT 1", None, "CSV") + .run_query_bearer("svc-1", "SELECT 1", None, "CSV", false) .await .expect_err("expected AuthMismatch"); assert!( @@ -115,13 +120,83 @@ async fn run_query_bearer_on_basic_auth_client_is_auth_mismatch() { ); } +#[tokio::test] +async fn run_query_wake_service_sends_wake_header() { + let mock = start_mock_query_host(200, "1\n").await; + let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); + + client + .run_query_bearer("svc-1", "SELECT 1", None, "TabSeparated", true) + .await + .expect("run_query_bearer failed"); + + let requests = mock.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].headers.get("wake-service").unwrap(), "true"); +} + +// ── 206 responses: the query host's service-state protocol ───────────────── +// +// An idled or stopped service answers 206 with `{"data": ""}` instead +// of running the query. `Confirm wake service` invites a resend with the +// `wake-service: true` header (which wakes the service); `Service is +// stopped` is terminal until the service is started. + +#[tokio::test] +async fn run_query_206_confirm_wake_maps_to_service_idle() { + let mock = start_mock_query_host(206, r#"{"data":"Confirm wake service"}"#).await; + let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); + + let err = client + .run_query_bearer("svc-1", "SELECT 1", None, "CSV", false) + .await + .expect_err("expected ServiceIdle"); + assert!( + matches!(err, Error::ServiceIdle), + "expected ServiceIdle, got: {err:?}" + ); +} + +#[tokio::test] +async fn run_query_206_service_stopped_maps_to_service_stopped() { + let mock = start_mock_query_host(206, r#"{"data":"Service is stopped"}"#).await; + let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); + + let err = client + .run_query_bearer("svc-1", "SELECT 1", None, "CSV", false) + .await + .expect_err("expected ServiceStopped"); + assert!( + matches!(err, Error::ServiceStopped), + "expected ServiceStopped, got: {err:?}" + ); +} + +#[tokio::test] +async fn run_query_206_unrecognized_body_maps_to_api_error() { + let mock = start_mock_query_host(206, r#"{"data":"Something new"}"#).await; + let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); + + let err = client + .run_query_bearer("svc-1", "SELECT 1", None, "CSV", false) + .await + .expect_err("expected Api error"); + match err { + Error::Api { status, message } => { + assert_eq!(status, 206); + assert_eq!(message, r#"{"data":"Something new"}"#); + } + other => panic!("expected Error::Api, got: {other:?}"), + } +} + #[tokio::test] async fn run_query_non_success_status_maps_to_api_error() { let mock = start_mock_query_host(404, "query endpoint not found").await; let client = Client::with_bearer_token(mock.uri(), "oauth-token").with_query_host(mock.uri()); let err = client - .run_query_bearer("svc-1", "SELECT 1", None, "CSV") + .run_query_bearer("svc-1", "SELECT 1", None, "CSV", false) .await .expect_err("expected Api error"); match err { diff --git a/crates/clickhousectl/src/cloud/commands.rs b/crates/clickhousectl/src/cloud/commands.rs index 489c3b2..ee3ebc1 100644 --- a/crates/clickhousectl/src/cloud/commands.rs +++ b/crates/clickhousectl/src/cloud/commands.rs @@ -2084,11 +2084,19 @@ pub async fn service_query( // query-endpoint configuration needed on the service. // `--no-auto-enable` is a no-op here since nothing is ever // provisioned. - client - .api() - .run_query_bearer(&service_id, &sql, opts.database.as_deref(), &format) - .await - .map_err(|e| client.convert_error(e))? + let run = |wake: bool| { + client + .api() + .run_query_bearer(&service_id, &sql, opts.database.as_deref(), &format, wake) + }; + let result = match run(false).await { + Err(clickhouse_cloud_api::Error::ServiceIdle) => { + eprint_waking_service(&service.name); + run(true).await + } + other => other, + }; + result.map_err(|e| convert_query_error(client, e, &service.name))? } else { let key = match credentials::get_service_query_key(&service_id) { Some(k) => k, @@ -2113,18 +2121,28 @@ pub async fn service_query( } }; - client - .api() - .run_query( + // The query host normally wakes an idled service on its own for + // Query API key auth, but handle the wake confirmation here too so + // both auth paths behave the same if it ever asks. + let run = |wake: bool| { + client.api().run_query( &service_id, &key.key_id, &key.key_secret, &sql, opts.database.as_deref(), &format, + wake, ) - .await - .map_err(|e| client.convert_error(e))? + }; + let result = match run(false).await { + Err(clickhouse_cloud_api::Error::ServiceIdle) => { + eprint_waking_service(&service.name); + run(true).await + } + other => other, + }; + result.map_err(|e| convert_query_error(client, e, &service.name))? }; use futures_util::StreamExt; @@ -2142,6 +2160,29 @@ pub async fn service_query( Ok(()) } +/// Stderr notice shown when the query host reports the service is idled and +/// the CLI resends the query with the wake confirmation. +fn eprint_waking_service(service_name: &str) { + eprintln!("Service '{service_name}' is idle; waking it (this may take a minute)..."); +} + +/// Map Query API errors to user-facing messages: a stopped service gets a +/// hint to start it (the query host never wakes a stopped service), the +/// rest go through the standard cloud error conversion. +fn convert_query_error( + client: &CloudClient, + err: clickhouse_cloud_api::Error, + service_name: &str, +) -> Box { + match err { + clickhouse_cloud_api::Error::ServiceStopped => format!( + "service '{service_name}' is stopped; start it with `clickhousectl cloud service start` and retry" + ) + .into(), + other => client.convert_error(other).into(), + } +} + fn read_query_sql( inline: Option<&str>, queries_file: Option<&str>, diff --git a/crates/clickhousectl/tests/cli_request_shape_test.rs b/crates/clickhousectl/tests/cli_request_shape_test.rs index 06fa29a..eae3890 100644 --- a/crates/clickhousectl/tests/cli_request_shape_test.rs +++ b/crates/clickhousectl/tests/cli_request_shape_test.rs @@ -18,7 +18,7 @@ use std::path::PathBuf; use std::process::Command; use serde_json::Value; -use wiremock::matchers::{method, path, path_regex}; +use wiremock::matchers::{header, method, path, path_regex}; use wiremock::{Mock, MockServer, ResponseTemplate}; /// Locate the `clickhousectl` binary. cargo populates `CARGO_BIN_EXE_` @@ -1457,6 +1457,159 @@ async fn service_query_with_stored_key_sends_basic_auth_with_that_key() { ); } +// ── Idled / stopped services (query host 206 protocol) ───────────────────── +// +// An idled or stopped service answers the run request with 206 and +// `{"data": ""}` instead of executing the query. For `Confirm wake +// service` the CLI must resend the query once with the `wake-service: true` +// header (waking the service, like the SQL console does after prompting); +// for `Service is stopped` it must fail with a hint to start the service. + +/// Write an OAuth tokens.json into `ch_dir` so the binary authenticates with +/// a bearer token against the given control plane. +fn write_oauth_tokens(ch_dir: &std::path::Path, control_uri: &str) { + let tokens = serde_json::json!({ + "access_token": "test-bearer-token", + "refresh_token": "unused", + "expires_at": 4102444800u64, // 2100-01-01: never expires in tests + "api_url": format!("{control_uri}/v1"), + }); + std::fs::write( + ch_dir.join("tokens.json"), + serde_json::to_vec(&tokens).unwrap(), + ) + .unwrap(); +} + +#[tokio::test] +async fn service_query_resends_with_wake_header_when_service_is_idle() { + let control = start_mock_control_plane_with_service().await; + + // Query host that refuses attempts without the wake confirmation: the + // header-matched mock (higher priority) runs the query, the fallback + // answers 206 `Confirm wake service`. + let query_host = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/service/{QUERY_TEST_SERVICE_ID}/run"))) + .and(header("wake-service", "true")) + .respond_with(ResponseTemplate::new(200).set_body_string("1\n")) + .with_priority(1) + .mount(&query_host) + .await; + Mock::given(method("POST")) + .and(path(format!("/service/{QUERY_TEST_SERVICE_ID}/run"))) + .respond_with( + ResponseTemplate::new(206).set_body_string(r#"{"data":"Confirm wake service"}"#), + ) + .with_priority(5) + .mount(&query_host) + .await; + + let dir = tempfile::tempdir().unwrap(); + let ch_dir = dir.path().join(".clickhouse"); + std::fs::create_dir_all(&ch_dir).unwrap(); + write_oauth_tokens(&ch_dir, &control.uri()); + + let url = control.uri(); + let output = Command::new(clickhousectl_binary()) + .args([ + "cloud", + "--url", + &url, + "service", + "query", + "--id", + QUERY_TEST_SERVICE_ID, + "--org-id", + "org-1", + "--query", + "SELECT 1", + ]) + .current_dir(dir.path()) + .env_remove("CLICKHOUSE_CLOUD_API_KEY") + .env_remove("CLICKHOUSE_CLOUD_API_SECRET") + .env("CLICKHOUSE_CLOUD_QUERY_HOST", query_host.uri()) + .output() + .expect("failed to spawn clickhousectl"); + assert_success(&output); + + // Exactly two attempts: the refused one without the wake header, then + // the resend carrying the wake confirmation. + let query_requests = query_host.received_requests().await.unwrap(); + assert_eq!(query_requests.len(), 2); + assert!( + query_requests[0].headers.get("wake-service").is_none(), + "first attempt must not pre-emptively wake the service", + ); + assert_eq!( + query_requests[1].headers.get("wake-service").unwrap(), + "true" + ); + + // The 206 body must not leak into the query output, and the user is + // told about the wake on stderr. + assert_eq!(String::from_utf8_lossy(&output.stdout), "1\n"); + assert!( + String::from_utf8_lossy(&output.stderr).contains("idle"), + "stderr should mention the service is idle:\n{}", + String::from_utf8_lossy(&output.stderr), + ); +} + +#[tokio::test] +async fn service_query_fails_with_start_hint_when_service_is_stopped() { + let control = start_mock_control_plane_with_service().await; + + let query_host = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!("/service/{QUERY_TEST_SERVICE_ID}/run"))) + .respond_with(ResponseTemplate::new(206).set_body_string(r#"{"data":"Service is stopped"}"#)) + .mount(&query_host) + .await; + + let dir = tempfile::tempdir().unwrap(); + let ch_dir = dir.path().join(".clickhouse"); + std::fs::create_dir_all(&ch_dir).unwrap(); + write_oauth_tokens(&ch_dir, &control.uri()); + + let url = control.uri(); + let output = Command::new(clickhousectl_binary()) + .args([ + "cloud", + "--url", + &url, + "service", + "query", + "--id", + QUERY_TEST_SERVICE_ID, + "--org-id", + "org-1", + "--query", + "SELECT 1", + ]) + .current_dir(dir.path()) + .env_remove("CLICKHOUSE_CLOUD_API_KEY") + .env_remove("CLICKHOUSE_CLOUD_API_SECRET") + .env("CLICKHOUSE_CLOUD_QUERY_HOST", query_host.uri()) + .output() + .expect("failed to spawn clickhousectl"); + + assert!( + !output.status.success(), + "querying a stopped service must fail\nstdout:\n{}", + String::from_utf8_lossy(&output.stdout), + ); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("stopped") && stderr.contains("service start"), + "stderr should say the service is stopped and hint at `service start`:\n{stderr}", + ); + + // A stopped service is never woken: no wake-service resend. + let query_requests = query_host.received_requests().await.unwrap(); + assert_eq!(query_requests.len(), 1); +} + // Shell env vars must win over `.env` — if both are set, the request is // signed with the shell values, never the file values.