cloud service query: support OAuth login without Query API key provisioning

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 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)
 
@@ -495,17 +495,20 @@ clickhousectl cloud service delete <service-id> --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.<service-id>`, 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.<service-id>`, 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`): 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.
 
 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.
+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`.
 
-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.[control-plane.]<domain>` → `queries.<domain>`, e.g. `https://queries.clickhouse.cloud` for production). Set `CLICKHOUSE_CLOUD_QUERY_HOST` to override it.
 
 ### Postgres (beta)
 

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,31 @@ pub struct Client {
     http: reqwest::Client,
     base_url: String,
     auth: Auth,
+    /// Explicit Query API host override; see [`Client::with_query_host`].
+    query_host: Option<String>,
+}
+
+/// 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. 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<String> {
+    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);
+    let port = parsed
+        .port()
+        .map(|p| format!(":{p}"))
+        .unwrap_or_default();
+    Some(format!("{}://queries.{}{}", parsed.scheme(), rest, port))
 }
 
 impl Client {
@@ -41,6 +78,7 @@ impl Client {
                 key_id: key_id.into(),
                 key_secret: key_secret.into(),
             },
+            query_host: None,
         }
     }
 
@@ -55,6 +93,7 @@ impl Client {
             auth: Auth::Bearer {
                 token: token.into(),
             },
+            query_host: None,
         }
     }
 
@@ -75,6 +114,7 @@ impl Client {
                 key_id: key_id.into(),
                 key_secret: key_secret.into(),
             },
+            query_host: None,
         }
     }
 
@@ -93,6 +133,7 @@ impl Client {
             auth: Auth::Bearer {
                 token: token.into(),
             },
+            query_host: None,
         }
     }
 
@@ -112,6 +153,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.<domain>` → `queries.<domain>`), falling back to the production
+    /// host `https://queries.clickhouse.cloud`.
+    pub fn with_query_host(mut self, host: impl Into<String>) -> 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,15 +190,19 @@ 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.
+    ///
+    /// `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,
@@ -141,6 +211,66 @@ impl Client {
         sql: &str,
         database: Option<&str>,
         format: &str,
+        wake_service: bool,
+    ) -> Result<reqwest::Response, Error> {
+        self.run_query_with(
+            QueryAuth::Basic { key_id, key_secret },
+            service_id,
+            sql,
+            database,
+            format,
+            wake_service,
+        )
+        .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 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.
+    ///
+    /// `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,
+        service_id: &str,
+        sql: &str,
+        database: Option<&str>,
+        format: &str,
+        wake_service: bool,
+    ) -> Result<reqwest::Response, Error> {
+        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,
+            wake_service,
+        )
+        .await
+    }
+
+    async fn run_query_with(
+        &self,
+        auth: QueryAuth<'_>,
+        service_id: &str,
+        sql: &str,
+        database: Option<&str>,
+        format: &str,
+        wake_service: bool,
     ) -> Result<reqwest::Response, Error> {
         #[derive(serde::Serialize)]
         #[serde(rename_all = "camelCase")]
@@ -151,11 +281,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,19 +293,55 @@ 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");
+        // `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.
+        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();
+        // 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<String>,
+            }
+            let data = serde_json::from_str::<StateBody>(&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 {
@@ -2702,4 +2866,66 @@ 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.control-plane.clickhouse-staging.com").as_deref(),
+            Some("https://queries.clickhouse-staging.com")
+        );
+    }
+
+    #[test]
+    fn derive_query_host_dev() {
+        assert_eq!(
+            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);
+        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);
+    }
+
+    #[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")
+        );
+    }
+}