Skip to main content

mongreldb_client/
lib.rs

1//! mongreldb-client — a lightweight HTTP client for `mongreldb-server`.
2//!
3//! Hardened surface:
4//! - Every call checks the HTTP status and returns a typed [`ClientError`] on
5//!   non-2xx (the server's `KitErrorEnvelope` is decoded into
6//!   [`KitError::unique_violation`]/`fk_violation`/etc.).
7//! - Typed [`KitTxnRequest`]/[`KitTxnResponse`] models mirror the server's
8//!   `/kit/txn` batch endpoint.
9//! - [`MongrelClient::kit_schema`] fetches table metadata (columns +
10//!   constraints).
11//! - Legacy SQL read ([`MongrelClient::sql`]) returns Arrow IPC batches.
12
13use std::io::{Cursor, Read};
14
15use arrow::ipc::reader::FileReader;
16use arrow::record_batch::RecordBatch;
17use secrecy::ExposeSecret;
18pub use secrecy::SecretString;
19use serde::{Deserialize, Serialize};
20use thiserror::Error;
21use zeroize::Zeroizing;
22
23/// A typed client error. Network/IO failures and non-2xx HTTP responses both
24/// surface here; constraint violations from `/kit/txn` are decoded into the
25/// matching variant so callers can branch on them.
26#[derive(Debug, Error)]
27pub enum ClientError {
28    #[error("io/transport error: {0}")]
29    Transport(String),
30    #[error("http {status}: {body}")]
31    Http { status: u16, body: String },
32    #[error("decode error: {0}")]
33    Decode(String),
34    #[error("kit error: {code}: {message}")]
35    Kit {
36        code: KitErrorCode,
37        message: String,
38        op_index: Option<usize>,
39        status: u16,
40        committed: Option<bool>,
41        epoch: Option<u64>,
42        epoch_text: Option<String>,
43        retryable: Option<bool>,
44    },
45    #[error("remote query {code}: {message}")]
46    Query {
47        status: u16,
48        code: RemoteQueryErrorCode,
49        message: String,
50        response: Box<RemoteQueryErrorResponse>,
51    },
52    #[error("query {query_id} outcome unknown after transport loss: {message}")]
53    QueryOutcomeUnknown {
54        query_id: String,
55        message: String,
56        status: Option<Box<RemoteQueryStatus>>,
57        cancel_outcome: Option<RemoteCancelOutcome>,
58    },
59}
60
61/// Typed error codes mirrored from the server's Kit error envelope.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum KitErrorCode {
64    UniqueViolation,
65    FkViolation,
66    CheckViolation,
67    ProcedureNotFound,
68    ProcedureValidation,
69    ProcedureExecution,
70    TriggerNotFound,
71    TriggerValidation,
72    Conflict,
73    AuthRequired,
74    PermissionDenied,
75    DeadlineExceeded,
76    WorkBudgetExceeded,
77    Cancelled,
78    CommitOutcome,
79    QueryOutcomeUnknown,
80    IdempotencyKeyReuseMismatch,
81    IdempotencyStoreFull,
82    IdempotencyStoreUnavailable,
83    InvalidIdempotencyKey,
84    BadRequest,
85    NotFound,
86    Internal,
87    Other,
88}
89
90impl std::fmt::Display for KitErrorCode {
91    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92        f.write_str(self.as_str())
93    }
94}
95
96impl KitErrorCode {
97    fn as_str(&self) -> &'static str {
98        match self {
99            KitErrorCode::UniqueViolation => "UNIQUE_VIOLATION",
100            KitErrorCode::FkViolation => "FK_VIOLATION",
101            KitErrorCode::CheckViolation => "CHECK_VIOLATION",
102            KitErrorCode::ProcedureNotFound => "PROCEDURE_NOT_FOUND",
103            KitErrorCode::ProcedureValidation => "PROCEDURE_VALIDATION",
104            KitErrorCode::ProcedureExecution => "PROCEDURE_EXECUTION",
105            KitErrorCode::TriggerNotFound => "TRIGGER_NOT_FOUND",
106            KitErrorCode::TriggerValidation => "TRIGGER_VALIDATION",
107            KitErrorCode::Conflict => "CONFLICT",
108            KitErrorCode::AuthRequired => "AUTH_REQUIRED",
109            KitErrorCode::PermissionDenied => "PERMISSION_DENIED",
110            KitErrorCode::DeadlineExceeded => "DEADLINE_EXCEEDED",
111            KitErrorCode::WorkBudgetExceeded => "WORK_BUDGET_EXCEEDED",
112            KitErrorCode::Cancelled => "CANCELLED",
113            KitErrorCode::CommitOutcome => "COMMIT_OUTCOME",
114            KitErrorCode::QueryOutcomeUnknown => "QUERY_OUTCOME_UNKNOWN",
115            KitErrorCode::IdempotencyKeyReuseMismatch => "IDEMPOTENCY_KEY_REUSE_MISMATCH",
116            KitErrorCode::IdempotencyStoreFull => "IDEMPOTENCY_STORE_FULL",
117            KitErrorCode::IdempotencyStoreUnavailable => "IDEMPOTENCY_STORE_UNAVAILABLE",
118            KitErrorCode::InvalidIdempotencyKey => "INVALID_IDEMPOTENCY_KEY",
119            KitErrorCode::BadRequest => "BAD_REQUEST",
120            KitErrorCode::NotFound => "NOT_FOUND",
121            KitErrorCode::Internal => "INTERNAL",
122            KitErrorCode::Other => "OTHER",
123        }
124    }
125
126    fn from_str(s: &str) -> Self {
127        match s {
128            "UNIQUE_VIOLATION" => KitErrorCode::UniqueViolation,
129            "FK_VIOLATION" => KitErrorCode::FkViolation,
130            "CHECK_VIOLATION" => KitErrorCode::CheckViolation,
131            "PROCEDURE_NOT_FOUND" => KitErrorCode::ProcedureNotFound,
132            "PROCEDURE_VALIDATION" => KitErrorCode::ProcedureValidation,
133            "PROCEDURE_EXECUTION" => KitErrorCode::ProcedureExecution,
134            "TRIGGER_NOT_FOUND" => KitErrorCode::TriggerNotFound,
135            "TRIGGER_VALIDATION" => KitErrorCode::TriggerValidation,
136            "CONFLICT" => KitErrorCode::Conflict,
137            "AUTH_REQUIRED" => KitErrorCode::AuthRequired,
138            "PERMISSION_DENIED" => KitErrorCode::PermissionDenied,
139            "DEADLINE_EXCEEDED" => KitErrorCode::DeadlineExceeded,
140            "WORK_BUDGET_EXCEEDED" => KitErrorCode::WorkBudgetExceeded,
141            "CANCELLED" => KitErrorCode::Cancelled,
142            "COMMIT_OUTCOME" => KitErrorCode::CommitOutcome,
143            "QUERY_OUTCOME_UNKNOWN" => KitErrorCode::QueryOutcomeUnknown,
144            "IDEMPOTENCY_KEY_REUSE_MISMATCH" => KitErrorCode::IdempotencyKeyReuseMismatch,
145            "IDEMPOTENCY_STORE_FULL" => KitErrorCode::IdempotencyStoreFull,
146            "IDEMPOTENCY_STORE_UNAVAILABLE" => KitErrorCode::IdempotencyStoreUnavailable,
147            "INVALID_IDEMPOTENCY_KEY" => KitErrorCode::InvalidIdempotencyKey,
148            "BAD_REQUEST" => KitErrorCode::BadRequest,
149            "NOT_FOUND" => KitErrorCode::NotFound,
150            "INTERNAL" => KitErrorCode::Internal,
151            _ => KitErrorCode::Other,
152        }
153    }
154}
155
156impl From<reqwest::Error> for ClientError {
157    fn from(e: reqwest::Error) -> Self {
158        ClientError::Transport(e.to_string())
159    }
160}
161
162impl From<std::io::Error> for ClientError {
163    fn from(e: std::io::Error) -> Self {
164        ClientError::Transport(e.to_string())
165    }
166}
167
168pub type ClientResult<T> = std::result::Result<T, ClientError>;
169
170const MAX_CONTROL_RESPONSE_BYTES: u64 = 1024 * 1024;
171// Server output-byte limits cover result data, not the enclosing JSON/Arrow
172// protocol metadata. Keep bounded headroom for a valid 64 MiB result.
173const MAX_SQL_RESPONSE_BYTES: u64 = 65 * 1024 * 1024;
174const MAX_REPLICATION_WAL_RESPONSE_BYTES: u64 = 256 * 1024 * 1024;
175const MAX_REPLICATION_SNAPSHOT_BYTES: u64 = 512 * 1024 * 1024;
176
177struct StrictJsonValue(serde_json::Value);
178
179impl<'de> Deserialize<'de> for StrictJsonValue {
180    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
181    where
182        D: serde::Deserializer<'de>,
183    {
184        struct Visitor;
185
186        impl<'de> serde::de::Visitor<'de> for Visitor {
187            type Value = StrictJsonValue;
188
189            fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
190                formatter.write_str("a JSON value without duplicate object keys")
191            }
192
193            fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
194                Ok(StrictJsonValue(value.into()))
195            }
196
197            fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
198                Ok(StrictJsonValue(value.into()))
199            }
200
201            fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
202                Ok(StrictJsonValue(value.into()))
203            }
204
205            fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
206            where
207                E: serde::de::Error,
208            {
209                serde_json::Number::from_f64(value)
210                    .map(serde_json::Value::Number)
211                    .map(StrictJsonValue)
212                    .ok_or_else(|| E::custom("non-finite JSON number"))
213            }
214
215            fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
216                Ok(StrictJsonValue(value.into()))
217            }
218
219            fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
220                Ok(StrictJsonValue(value.into()))
221            }
222
223            fn visit_none<E>(self) -> Result<Self::Value, E> {
224                Ok(StrictJsonValue(serde_json::Value::Null))
225            }
226
227            fn visit_unit<E>(self) -> Result<Self::Value, E> {
228                Ok(StrictJsonValue(serde_json::Value::Null))
229            }
230
231            fn visit_some<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
232            where
233                D: serde::Deserializer<'de>,
234            {
235                <StrictJsonValue as Deserialize>::deserialize(deserializer)
236            }
237
238            fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
239            where
240                A: serde::de::SeqAccess<'de>,
241            {
242                let mut values = Vec::with_capacity(sequence.size_hint().unwrap_or(0));
243                while let Some(value) = sequence.next_element::<StrictJsonValue>()? {
244                    values.push(value.0);
245                }
246                Ok(StrictJsonValue(serde_json::Value::Array(values)))
247            }
248
249            fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
250            where
251                A: serde::de::MapAccess<'de>,
252            {
253                let mut values = serde_json::Map::with_capacity(map.size_hint().unwrap_or(0));
254                while let Some(key) = map.next_key::<String>()? {
255                    if values.contains_key(&key) {
256                        return Err(serde::de::Error::custom(format!(
257                            "duplicate JSON object key {key:?}"
258                        )));
259                    }
260                    let value = map.next_value::<StrictJsonValue>()?;
261                    values.insert(key, value.0);
262                }
263                Ok(StrictJsonValue(serde_json::Value::Object(values)))
264            }
265        }
266
267        deserializer.deserialize_any(Visitor)
268    }
269}
270
271fn strict_json_value(bytes: &[u8]) -> Result<serde_json::Value, String> {
272    let mut deserializer = serde_json::Deserializer::from_slice(bytes);
273    let value =
274        StrictJsonValue::deserialize(&mut deserializer).map_err(|error| error.to_string())?;
275    deserializer.end().map_err(|error| error.to_string())?;
276    Ok(value.0)
277}
278
279fn strict_json<T: serde::de::DeserializeOwned>(bytes: &[u8], context: &str) -> ClientResult<T> {
280    let value = strict_json_value(bytes)
281        .map_err(|error| ClientError::Decode(format!("invalid {context} response: {error}")))?;
282    serde_json::from_value(value)
283        .map_err(|error| ClientError::Decode(format!("invalid {context} response: {error}")))
284}
285
286fn strict_roundtrip_json<T>(bytes: &[u8], context: &str) -> ClientResult<T>
287where
288    T: serde::de::DeserializeOwned + Serialize,
289{
290    let value = strict_json_value(bytes)
291        .map_err(|error| ClientError::Decode(format!("invalid {context}: {error}")))?;
292    let parsed: T = serde_json::from_value(value.clone())
293        .map_err(|error| ClientError::Decode(format!("invalid {context}: {error}")))?;
294    if serde_json::to_value(&parsed)
295        .map_err(|error| ClientError::Decode(format!("invalid {context}: {error}")))?
296        != value
297    {
298        return Err(ClientError::Decode(format!(
299            "invalid {context}: unknown or non-canonical fields"
300        )));
301    }
302    Ok(parsed)
303}
304
305fn bounded_blocking_bytes(
306    mut response: reqwest::blocking::Response,
307    limit: u64,
308) -> Result<Vec<u8>, String> {
309    if response
310        .content_length()
311        .is_some_and(|length| length > limit)
312    {
313        return Err(format!("response exceeds {limit} bytes"));
314    }
315    let mut bytes = Vec::new();
316    response
317        .by_ref()
318        .take(limit.saturating_add(1))
319        .read_to_end(&mut bytes)
320        .map_err(|error| error.to_string())?;
321    if bytes.len() as u64 > limit {
322        return Err(format!("response exceeds {limit} bytes"));
323    }
324    Ok(bytes)
325}
326
327async fn bounded_async_bytes(
328    mut response: reqwest::Response,
329    limit: u64,
330) -> Result<Vec<u8>, String> {
331    if response
332        .content_length()
333        .is_some_and(|length| length > limit)
334    {
335        return Err(format!("response exceeds {limit} bytes"));
336    }
337    let mut bytes = Vec::new();
338    while let Some(chunk) = response.chunk().await.map_err(|error| error.to_string())? {
339        if (bytes.len() as u64).saturating_add(chunk.len() as u64) > limit {
340            return Err(format!("response exceeds {limit} bytes"));
341        }
342        bytes.extend_from_slice(&chunk);
343    }
344    Ok(bytes)
345}
346
347fn decode_blocking_json<T: serde::de::DeserializeOwned>(
348    response: reqwest::blocking::Response,
349    limit: u64,
350    context: &str,
351) -> ClientResult<T> {
352    let bytes = bounded_blocking_bytes(response, limit)
353        .map_err(|error| ClientError::Decode(format!("invalid {context} response: {error}")))?;
354    strict_json(&bytes, context)
355}
356
357async fn decode_async_json<T: serde::de::DeserializeOwned>(
358    response: reqwest::Response,
359    limit: u64,
360    context: &str,
361) -> ClientResult<T> {
362    let bytes = bounded_async_bytes(response, limit)
363        .await
364        .map_err(|error| ClientError::Decode(format!("invalid {context} response: {error}")))?;
365    strict_json(&bytes, context)
366}
367
368#[derive(Clone)]
369pub enum RemoteAuth {
370    Bearer(SecretString),
371    Basic {
372        username: String,
373        password: SecretString,
374    },
375}
376
377impl std::fmt::Debug for RemoteAuth {
378    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
379        match self {
380            Self::Bearer(_) => formatter.write_str("Bearer([REDACTED])"),
381            Self::Basic { username, .. } => formatter
382                .debug_struct("Basic")
383                .field("username", username)
384                .field("password", &"[REDACTED]")
385                .finish(),
386        }
387    }
388}
389
390#[derive(Debug, Clone, Default)]
391pub struct RemoteOptions {
392    pub auth: Option<RemoteAuth>,
393    pub transport_timeout: Option<std::time::Duration>,
394}
395
396#[derive(Clone)]
397pub struct MongrelClient {
398    base_url: String,
399    client: reqwest::blocking::Client,
400    transport: ClientTransportOptions,
401}
402
403#[derive(Clone, Copy, Default)]
404struct ClientTransportOptions {
405    connect_timeout: Option<std::time::Duration>,
406    request_timeout: Option<std::time::Duration>,
407    pool_idle_timeout: Option<std::time::Duration>,
408}
409
410/// Builder for a blocking [`MongrelClient`]. Authorization is installed as a
411/// sensitive default header, so every route uses it and debug output redacts it.
412pub struct MongrelClientBuilder {
413    base_url: String,
414    invalid_base_url: bool,
415    authorization: Option<reqwest::header::HeaderValue>,
416    invalid_authorization: bool,
417    connect_timeout: Option<std::time::Duration>,
418    request_timeout: Option<std::time::Duration>,
419    pool_idle_timeout: Option<std::time::Duration>,
420}
421
422/// Builder for an [`AsyncMongrelClient`].
423pub struct AsyncMongrelClientBuilder {
424    base_url: String,
425    invalid_base_url: bool,
426    authorization: Option<reqwest::header::HeaderValue>,
427    invalid_authorization: bool,
428    connect_timeout: Option<std::time::Duration>,
429    request_timeout: Option<std::time::Duration>,
430    pool_idle_timeout: Option<std::time::Duration>,
431}
432
433fn bearer_header(token: &str) -> ClientResult<reqwest::header::HeaderValue> {
434    let value = Zeroizing::new(format!("Bearer {token}"));
435    sensitive_header(value.as_str())
436}
437
438fn basic_header(username: &str, password: &str) -> ClientResult<reqwest::header::HeaderValue> {
439    let credentials = Zeroizing::new(format!("{username}:{password}"));
440    let encoded = Zeroizing::new(base64::Engine::encode(
441        &base64::engine::general_purpose::STANDARD,
442        credentials.as_bytes(),
443    ));
444    let value = Zeroizing::new(format!("Basic {}", encoded.as_str()));
445    sensitive_header(value.as_str())
446}
447
448fn sensitive_header(value: &str) -> ClientResult<reqwest::header::HeaderValue> {
449    let mut value = reqwest::header::HeaderValue::from_str(value)
450        .map_err(|_| ClientError::Transport("invalid authorization credentials".into()))?;
451    value.set_sensitive(true);
452    Ok(value)
453}
454
455fn default_headers(
456    authorization: Option<reqwest::header::HeaderValue>,
457) -> reqwest::header::HeaderMap {
458    let mut headers = reqwest::header::HeaderMap::new();
459    if let Some(value) = authorization {
460        headers.insert(reqwest::header::AUTHORIZATION, value);
461    }
462    headers
463}
464
465fn url_with_segments(base_url: &str, segments: &[&str]) -> ClientResult<String> {
466    let mut url = reqwest::Url::parse(base_url)
467        .map_err(|_| ClientError::Transport("invalid client base URL".into()))?;
468    {
469        let mut path = url
470            .path_segments_mut()
471            .map_err(|_| ClientError::Transport("client base URL cannot contain a path".into()))?;
472        path.pop_if_empty();
473        path.extend(segments.iter().copied());
474    }
475    Ok(url.into())
476}
477
478fn blocking_http_client(
479    authorization: Option<reqwest::header::HeaderValue>,
480    transport: ClientTransportOptions,
481) -> ClientResult<reqwest::blocking::Client> {
482    let mut builder =
483        reqwest::blocking::Client::builder().default_headers(default_headers(authorization));
484    if let Some(timeout) = transport.connect_timeout {
485        builder = builder.connect_timeout(timeout);
486    }
487    if let Some(timeout) = transport.request_timeout {
488        builder = builder.timeout(timeout);
489    }
490    if let Some(timeout) = transport.pool_idle_timeout {
491        builder = builder.pool_idle_timeout(timeout);
492    }
493    Ok(builder.build()?)
494}
495
496fn async_http_client(
497    authorization: Option<reqwest::header::HeaderValue>,
498    transport: ClientTransportOptions,
499) -> ClientResult<reqwest::Client> {
500    let mut builder = reqwest::Client::builder().default_headers(default_headers(authorization));
501    if let Some(timeout) = transport.connect_timeout {
502        builder = builder.connect_timeout(timeout);
503    }
504    if let Some(timeout) = transport.request_timeout {
505        builder = builder.timeout(timeout);
506    }
507    if let Some(timeout) = transport.pool_idle_timeout {
508        builder = builder.pool_idle_timeout(timeout);
509    }
510    Ok(builder.build()?)
511}
512
513fn sanitized_base_url(value: &str) -> Option<String> {
514    let url = reqwest::Url::parse(value).ok()?;
515    if !matches!(url.scheme(), "http" | "https")
516        || url.host_str().is_none()
517        || !url.username().is_empty()
518        || url.password().is_some()
519        || url.query().is_some()
520        || url.fragment().is_some()
521    {
522        return None;
523    }
524    Some(url.as_str().trim_end_matches('/').to_owned())
525}
526
527impl MongrelClientBuilder {
528    pub fn connect_timeout(mut self, timeout: std::time::Duration) -> Self {
529        self.connect_timeout = Some(timeout);
530        self
531    }
532
533    pub fn request_timeout(mut self, timeout: std::time::Duration) -> Self {
534        self.request_timeout = Some(timeout);
535        self
536    }
537
538    pub fn pool_idle_timeout(mut self, timeout: std::time::Duration) -> Self {
539        self.pool_idle_timeout = Some(timeout);
540        self
541    }
542
543    pub fn bearer_token(mut self, token: impl AsRef<str>) -> Self {
544        match bearer_header(token.as_ref()) {
545            Ok(value) => self.authorization = Some(value),
546            Err(_) => self.invalid_authorization = true,
547        }
548        self
549    }
550
551    pub fn basic_auth(mut self, username: impl AsRef<str>, password: impl AsRef<str>) -> Self {
552        match basic_header(username.as_ref(), password.as_ref()) {
553            Ok(value) => self.authorization = Some(value),
554            Err(_) => self.invalid_authorization = true,
555        }
556        self
557    }
558
559    pub fn build(self) -> ClientResult<MongrelClient> {
560        if self.invalid_base_url {
561            return Err(ClientError::Transport(
562                "invalid MongrelDB server URL; use HTTP(S) without credentials, query, or fragment"
563                    .into(),
564            ));
565        }
566        if self.invalid_authorization {
567            return Err(ClientError::Transport(
568                "invalid authorization credentials".into(),
569            ));
570        }
571        let transport = ClientTransportOptions {
572            connect_timeout: self.connect_timeout,
573            request_timeout: self.request_timeout,
574            pool_idle_timeout: self.pool_idle_timeout,
575        };
576        let client = blocking_http_client(self.authorization, transport)?;
577        Ok(MongrelClient {
578            base_url: self.base_url,
579            client,
580            transport,
581        })
582    }
583}
584
585impl AsyncMongrelClientBuilder {
586    pub fn connect_timeout(mut self, timeout: std::time::Duration) -> Self {
587        self.connect_timeout = Some(timeout);
588        self
589    }
590
591    pub fn request_timeout(mut self, timeout: std::time::Duration) -> Self {
592        self.request_timeout = Some(timeout);
593        self
594    }
595
596    pub fn pool_idle_timeout(mut self, timeout: std::time::Duration) -> Self {
597        self.pool_idle_timeout = Some(timeout);
598        self
599    }
600
601    pub fn bearer_token(mut self, token: impl AsRef<str>) -> Self {
602        match bearer_header(token.as_ref()) {
603            Ok(value) => self.authorization = Some(value),
604            Err(_) => self.invalid_authorization = true,
605        }
606        self
607    }
608
609    pub fn basic_auth(mut self, username: impl AsRef<str>, password: impl AsRef<str>) -> Self {
610        match basic_header(username.as_ref(), password.as_ref()) {
611            Ok(value) => self.authorization = Some(value),
612            Err(_) => self.invalid_authorization = true,
613        }
614        self
615    }
616
617    pub fn build(self) -> ClientResult<AsyncMongrelClient> {
618        if self.invalid_base_url {
619            return Err(ClientError::Transport(
620                "invalid MongrelDB server URL; use HTTP(S) without credentials, query, or fragment"
621                    .into(),
622            ));
623        }
624        if self.invalid_authorization {
625            return Err(ClientError::Transport(
626                "invalid authorization credentials".into(),
627            ));
628        }
629        let transport = ClientTransportOptions {
630            connect_timeout: self.connect_timeout,
631            request_timeout: self.request_timeout,
632            pool_idle_timeout: self.pool_idle_timeout,
633        };
634        let client = async_http_client(self.authorization, transport)?;
635        Ok(AsyncMongrelClient {
636            base_url: self.base_url,
637            client,
638            transport,
639        })
640    }
641}
642
643impl MongrelClient {
644    pub fn try_with_bearer_token(mut self, token: impl AsRef<str>) -> ClientResult<Self> {
645        self.client = blocking_http_client(Some(bearer_header(token.as_ref())?), self.transport)?;
646        Ok(self)
647    }
648
649    pub fn with_bearer_token(self, token: impl AsRef<str>) -> ClientResult<Self> {
650        self.try_with_bearer_token(token)
651    }
652
653    pub fn try_with_basic_auth(
654        mut self,
655        username: impl AsRef<str>,
656        password: impl AsRef<str>,
657    ) -> ClientResult<Self> {
658        self.client = blocking_http_client(
659            Some(basic_header(username.as_ref(), password.as_ref())?),
660            self.transport,
661        )?;
662        Ok(self)
663    }
664
665    pub fn with_basic_auth(
666        self,
667        username: impl AsRef<str>,
668        password: impl AsRef<str>,
669    ) -> ClientResult<Self> {
670        self.try_with_basic_auth(username, password)
671    }
672}
673
674/// Async counterpart for Kit and health calls.
675#[derive(Clone)]
676pub struct AsyncMongrelClient {
677    base_url: String,
678    client: reqwest::Client,
679    transport: ClientTransportOptions,
680}
681
682const SQL_RECOVERY_WINDOW: std::time::Duration = std::time::Duration::from_secs(2);
683const SQL_RECOVERY_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(250);
684const SQL_RECOVERY_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(25);
685
686#[derive(Debug, Clone, Default)]
687pub struct SqlClientOptions {
688    pub query_id: Option<mongreldb_query::QueryId>,
689    pub timeout: Option<std::time::Duration>,
690}
691
692pub type RemoteSqlControlOptions = SqlClientOptions;
693
694#[derive(Debug, Clone)]
695pub struct SqlPageOptions {
696    pub query_id: Option<mongreldb_query::QueryId>,
697    pub timeout: Option<std::time::Duration>,
698    pub max_output_rows: Option<u64>,
699    pub max_output_bytes: Option<u64>,
700    pub page_size_rows: u64,
701    pub projection: Vec<String>,
702    pub max_page_bytes: Option<u64>,
703    pub max_page_tokens: Option<u64>,
704}
705
706impl SqlPageOptions {
707    pub fn new(page_size_rows: u64, projection: Vec<String>) -> Self {
708        Self {
709            query_id: None,
710            timeout: None,
711            max_output_rows: None,
712            max_output_bytes: None,
713            page_size_rows,
714            projection,
715            max_page_bytes: None,
716            max_page_tokens: None,
717        }
718    }
719}
720
721fn validate_sql_page_options(options: &SqlPageOptions) -> ClientResult<()> {
722    if options.page_size_rows == 0
723        || options.max_output_rows == Some(0)
724        || options.max_output_bytes == Some(0)
725        || options.max_page_bytes == Some(0)
726        || options.max_page_tokens == Some(0)
727    {
728        return Err(ClientError::Decode(
729            "SQL pagination row, byte, and token limits must be positive".into(),
730        ));
731    }
732    if options.projection.is_empty() || options.projection.len() > 128 {
733        return Err(ClientError::Decode(
734            "SQL pagination projection must contain 1 to 128 columns".into(),
735        ));
736    }
737    let mut seen = std::collections::HashSet::new();
738    let projection_bytes = options
739        .projection
740        .iter()
741        .map(String::len)
742        .fold(0usize, usize::saturating_add);
743    if projection_bytes > 16 * 1024
744        || options.projection.iter().any(|column| {
745            column.is_empty()
746                || column == "*"
747                || column.len() > 256
748                || !seen.insert(column.as_str())
749        })
750    {
751        return Err(ClientError::Decode(
752            "SQL pagination projection requires unique explicit column names of at most 256 bytes"
753                .into(),
754        ));
755    }
756    Ok(())
757}
758
759fn validate_remote_sql_page(
760    page: RemoteSqlPage,
761    initial_options: Option<&SqlPageOptions>,
762) -> Result<RemoteSqlPage, String> {
763    let metadata = &page.page;
764    let end = metadata
765        .offset
766        .checked_add(metadata.row_count)
767        .ok_or_else(|| "SQL page offset overflowed".to_owned())?;
768    if page.status != "completed" {
769        return Err("SQL page status is not completed".into());
770    }
771    if metadata.row_count != page.rows.len() {
772        return Err("SQL page row_count does not match rows".into());
773    }
774    if page.rows.iter().any(|row| !row.is_object()) {
775        return Err("SQL page rows must be JSON objects".into());
776    }
777    if metadata.projection.is_empty()
778        || metadata.projection.iter().any(|column| column.is_empty())
779        || metadata
780            .projection
781            .iter()
782            .collect::<std::collections::HashSet<_>>()
783            .len()
784            != metadata.projection.len()
785        || page.rows.iter().any(|row| match row.as_object() {
786            Some(object) => {
787                object.len() != metadata.projection.len()
788                    || metadata
789                        .projection
790                        .iter()
791                        .any(|column| !object.contains_key(column))
792            }
793            None => true,
794        })
795    {
796        return Err("SQL page rows do not exactly match the projection".into());
797    }
798    let byte_count = page.rows.iter().try_fold(2_usize, |bytes, row| {
799        serde_json::to_vec(row)
800            .map(|encoded| {
801                bytes
802                    .saturating_add(usize::from(bytes > 2))
803                    .saturating_add(encoded.len())
804            })
805            .map_err(|error| error.to_string())
806    })?;
807    if metadata.byte_count != byte_count
808        || metadata.estimated_tokens != byte_count.saturating_add(3) / 4
809    {
810        return Err("SQL page byte or token estimate is invalid".into());
811    }
812    if metadata.offset > metadata.total_rows || end > metadata.total_rows {
813        return Err("SQL page offset or row_count exceeds total_rows".into());
814    }
815    if metadata.limits.rows == 0
816        || metadata.limits.bytes == 0
817        || metadata.limits.tokens == 0
818        || metadata.row_count > metadata.limits.rows
819        || metadata.byte_count > metadata.limits.bytes
820        || metadata.estimated_tokens > metadata.limits.tokens
821    {
822        return Err("SQL page exceeds its declared limits".into());
823    }
824    if metadata.expires_at_ms == 0
825        || metadata.snapshot != "retained_result"
826        || metadata.token_estimate != "ceil(projected_json_bytes/4)"
827    {
828        return Err("SQL page metadata is invalid".into());
829    }
830    let has_more = end < metadata.total_rows;
831    if (has_more && metadata.row_count == 0)
832        || has_more != page.next_cursor.is_some()
833        || page
834            .next_cursor
835            .as_ref()
836            .is_some_and(|cursor| cursor.is_empty() || cursor.len() > 2_048)
837    {
838        return Err("SQL page continuation cursor is inconsistent".into());
839    }
840    if let Some(options) = initial_options {
841        if metadata.offset != 0
842            || metadata.projection != options.projection
843            || metadata.limits.rows as u64 > options.page_size_rows
844            || options
845                .max_page_bytes
846                .is_some_and(|limit| metadata.limits.bytes as u64 > limit)
847            || options
848                .max_page_tokens
849                .is_some_and(|limit| metadata.limits.tokens as u64 > limit)
850            || options
851                .max_output_rows
852                .is_some_and(|limit| metadata.total_rows as u64 > limit)
853            || options
854                .max_output_bytes
855                .is_some_and(|limit| metadata.byte_count as u64 > limit)
856        {
857            return Err("SQL page does not match the requested pagination options".into());
858        }
859    }
860    Ok(page)
861}
862
863#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
864#[serde(rename_all = "snake_case")]
865pub enum RemoteCancelOutcome {
866    Accepted,
867    AlreadyCancelling,
868    TooLate,
869    AlreadyFinished,
870    NotFound,
871    PreCancelled,
872}
873
874fn cancel_outcome_from_wire(value: Option<&str>) -> Option<RemoteCancelOutcome> {
875    match value {
876        Some("accepted" | "cancellation_requested") => Some(RemoteCancelOutcome::Accepted),
877        Some("already_cancelling" | "cancelling") => Some(RemoteCancelOutcome::AlreadyCancelling),
878        Some("too_late" | "commit_critical") => Some(RemoteCancelOutcome::TooLate),
879        Some("already_finished" | "finished") => Some(RemoteCancelOutcome::AlreadyFinished),
880        Some("pre_cancelled") => Some(RemoteCancelOutcome::PreCancelled),
881        Some("not_found") => Some(RemoteCancelOutcome::NotFound),
882        _ => None,
883    }
884}
885
886fn decode_cancel_outcome(
887    body: &serde_json::Value,
888    expected_query_id: mongreldb_query::QueryId,
889    http_status: u16,
890) -> ClientResult<RemoteCancelOutcome> {
891    const FIELDS: &[&str] = &[
892        "query_id",
893        "status",
894        "terminal_state",
895        "state",
896        "server_state",
897        "code",
898        "operation",
899        "committed",
900        "committed_statements",
901        "last_commit_epoch",
902        "last_commit_epoch_text",
903        "first_commit_statement_index",
904        "last_commit_statement_index",
905        "completed_statements",
906        "statement_index",
907        "cancel_outcome",
908        "cancellation_reason",
909        "retryable",
910        "outcome",
911        "terminal_error",
912        "error",
913        "trace",
914        "started_ms_ago",
915        "deadline_ms_remaining",
916        "session_id",
917    ];
918    let object = body
919        .as_object()
920        .ok_or_else(|| ClientError::Decode("cancellation response is not an object".into()))?;
921    if let Some(field) = object
922        .keys()
923        .find(|field| !FIELDS.contains(&field.as_str()))
924    {
925        return Err(ClientError::Decode(format!(
926            "cancellation response contains unknown field {field:?}"
927        )));
928    }
929    if let Some(outcome) = object.get("outcome") {
930        serde_json::from_value::<RemoteQueryOutcome>(outcome.clone()).map_err(|error| {
931            ClientError::Decode(format!("invalid cancellation outcome: {error}"))
932        })?;
933    }
934    if let Some(error) = object.get("error") {
935        serde_json::from_value::<RemoteQueryErrorBody>(error.clone())
936            .map_err(|error| ClientError::Decode(format!("invalid cancellation error: {error}")))?;
937    }
938    if let Some(error) = object
939        .get("terminal_error")
940        .filter(|value| !value.is_null())
941    {
942        serde_json::from_value::<RemoteTerminalError>(error.clone()).map_err(|error| {
943            ClientError::Decode(format!("invalid cancellation terminal error: {error}"))
944        })?;
945    }
946    if body.get("query_id").and_then(serde_json::Value::as_str)
947        != Some(expected_query_id.to_string().as_str())
948    {
949        return Err(ClientError::Decode(
950            "cancellation response query_id does not match the request".into(),
951        ));
952    }
953    let outcome = cancel_outcome_from_wire(
954        body.get("cancel_outcome")
955            .and_then(serde_json::Value::as_str),
956    );
957    let state = cancel_outcome_from_wire(body.get("state").and_then(serde_json::Value::as_str));
958    if outcome.is_some() && state.is_some() && outcome != state {
959        return Err(ClientError::Decode(
960            "cancellation response state and cancel_outcome disagree".into(),
961        ));
962    }
963    let outcome = outcome
964        .or(state)
965        .ok_or_else(|| ClientError::Decode("cancellation response has no valid outcome".into()))?;
966    let compatible = matches!(
967        (http_status, outcome),
968        (
969            202,
970            RemoteCancelOutcome::Accepted | RemoteCancelOutcome::PreCancelled
971        ) | (
972            200,
973            RemoteCancelOutcome::AlreadyCancelling | RemoteCancelOutcome::AlreadyFinished
974        ) | (409, RemoteCancelOutcome::TooLate)
975            | (404, RemoteCancelOutcome::NotFound)
976    );
977    if !compatible {
978        return Err(ClientError::Decode(
979            "cancellation HTTP status and outcome disagree".into(),
980        ));
981    }
982    Ok(outcome)
983}
984
985#[derive(Debug, Clone, PartialEq, Eq)]
986pub enum RemoteQueryErrorCode {
987    QueryCancelled,
988    DeadlineExceeded,
989    QueryIdConflict,
990    QueryRegistryFull,
991    CancelTooLate,
992    QueryAlreadyFinished,
993    QueryNotFound,
994    TransactionAborted,
995    NoSqlTransaction,
996    SavepointNotFound,
997    CommitOutcome,
998    QueryFailed,
999    QueryCancelledAfterCommit,
1000    DeadlineAfterCommit,
1001    ResultLimitExceeded,
1002    SerializationFailed,
1003    SerializationFailedAfterCommit,
1004    SerializationWorkerFailed,
1005    CapabilityUnsupported,
1006    QueryOutcomeUnknown,
1007    InvalidQueryOptions,
1008    IncompatibleSqlControls,
1009    InvalidIdempotencyKey,
1010    IdempotencyKeyReuseMismatch,
1011    IdempotencyRequiresJson,
1012    IdempotencyRequiresSingleWrite,
1013    IdempotencyStoreFull,
1014    IdempotencyUnsupportedInTransaction,
1015    IdempotencyStoreUnavailable,
1016    PaginationRequiresJson,
1017    PaginationRequiresSingleReadQuery,
1018    InvalidPaginationOptions,
1019    InvalidSqlCursor,
1020    SqlCursorExpired,
1021    SqlCursorNotFound,
1022    InvalidSqlProjection,
1023    InvalidPageOffset,
1024    SqlPageStoreFull,
1025    SqlAdmissionClosed,
1026    EntropyUnavailable,
1027    Other(String),
1028}
1029
1030impl RemoteQueryErrorCode {
1031    pub fn as_str(&self) -> &str {
1032        match self {
1033            Self::QueryCancelled => "QUERY_CANCELLED",
1034            Self::DeadlineExceeded => "DEADLINE_EXCEEDED",
1035            Self::QueryIdConflict => "QUERY_ID_CONFLICT",
1036            Self::QueryRegistryFull => "QUERY_REGISTRY_FULL",
1037            Self::CancelTooLate => "CANCEL_TOO_LATE",
1038            Self::QueryAlreadyFinished => "QUERY_ALREADY_FINISHED",
1039            Self::QueryNotFound => "QUERY_NOT_FOUND",
1040            Self::TransactionAborted => "TRANSACTION_ABORTED",
1041            Self::NoSqlTransaction => "NO_SQL_TRANSACTION",
1042            Self::SavepointNotFound => "SAVEPOINT_NOT_FOUND",
1043            Self::CommitOutcome => "COMMIT_OUTCOME",
1044            Self::QueryFailed => "QUERY_FAILED",
1045            Self::QueryCancelledAfterCommit => "QUERY_CANCELLED_AFTER_COMMIT",
1046            Self::DeadlineAfterCommit => "DEADLINE_AFTER_COMMIT",
1047            Self::ResultLimitExceeded => "RESULT_LIMIT_EXCEEDED",
1048            Self::SerializationFailed => "SERIALIZATION_FAILED",
1049            Self::SerializationFailedAfterCommit => "SERIALIZATION_FAILED_AFTER_COMMIT",
1050            Self::SerializationWorkerFailed => "SERIALIZATION_WORKER_FAILED",
1051            Self::CapabilityUnsupported => "CAPABILITY_UNSUPPORTED",
1052            Self::QueryOutcomeUnknown => "QUERY_OUTCOME_UNKNOWN",
1053            Self::InvalidQueryOptions => "INVALID_QUERY_OPTIONS",
1054            Self::IncompatibleSqlControls => "INCOMPATIBLE_SQL_CONTROLS",
1055            Self::InvalidIdempotencyKey => "INVALID_IDEMPOTENCY_KEY",
1056            Self::IdempotencyKeyReuseMismatch => "IDEMPOTENCY_KEY_REUSE_MISMATCH",
1057            Self::IdempotencyRequiresJson => "IDEMPOTENCY_REQUIRES_JSON",
1058            Self::IdempotencyRequiresSingleWrite => "IDEMPOTENCY_REQUIRES_SINGLE_WRITE",
1059            Self::IdempotencyStoreFull => "IDEMPOTENCY_STORE_FULL",
1060            Self::IdempotencyUnsupportedInTransaction => "IDEMPOTENCY_UNSUPPORTED_IN_TRANSACTION",
1061            Self::IdempotencyStoreUnavailable => "IDEMPOTENCY_STORE_UNAVAILABLE",
1062            Self::PaginationRequiresJson => "PAGINATION_REQUIRES_JSON",
1063            Self::PaginationRequiresSingleReadQuery => "PAGINATION_REQUIRES_SINGLE_READ_QUERY",
1064            Self::InvalidPaginationOptions => "INVALID_PAGINATION_OPTIONS",
1065            Self::InvalidSqlCursor => "INVALID_SQL_CURSOR",
1066            Self::SqlCursorExpired => "SQL_CURSOR_EXPIRED",
1067            Self::SqlCursorNotFound => "SQL_CURSOR_NOT_FOUND",
1068            Self::InvalidSqlProjection => "INVALID_SQL_PROJECTION",
1069            Self::InvalidPageOffset => "INVALID_PAGE_OFFSET",
1070            Self::SqlPageStoreFull => "SQL_PAGE_STORE_FULL",
1071            Self::SqlAdmissionClosed => "SQL_ADMISSION_CLOSED",
1072            Self::EntropyUnavailable => "ENTROPY_UNAVAILABLE",
1073            Self::Other(code) => code,
1074        }
1075    }
1076
1077    fn from_code(code: &str) -> Self {
1078        match code {
1079            "QUERY_CANCELLED" => Self::QueryCancelled,
1080            "DEADLINE_EXCEEDED" => Self::DeadlineExceeded,
1081            "QUERY_ID_CONFLICT" => Self::QueryIdConflict,
1082            "QUERY_REGISTRY_FULL" => Self::QueryRegistryFull,
1083            "CANCEL_TOO_LATE" => Self::CancelTooLate,
1084            "QUERY_ALREADY_FINISHED" => Self::QueryAlreadyFinished,
1085            "QUERY_NOT_FOUND" => Self::QueryNotFound,
1086            "TRANSACTION_ABORTED" => Self::TransactionAborted,
1087            "NO_SQL_TRANSACTION" => Self::NoSqlTransaction,
1088            "SAVEPOINT_NOT_FOUND" => Self::SavepointNotFound,
1089            "COMMIT_OUTCOME" => Self::CommitOutcome,
1090            "QUERY_FAILED" => Self::QueryFailed,
1091            "QUERY_CANCELLED_AFTER_COMMIT" => Self::QueryCancelledAfterCommit,
1092            "DEADLINE_AFTER_COMMIT" => Self::DeadlineAfterCommit,
1093            "RESULT_LIMIT_EXCEEDED" => Self::ResultLimitExceeded,
1094            "SERIALIZATION_FAILED" => Self::SerializationFailed,
1095            "SERIALIZATION_FAILED_AFTER_COMMIT" => Self::SerializationFailedAfterCommit,
1096            "SERIALIZATION_WORKER_FAILED" => Self::SerializationWorkerFailed,
1097            "CAPABILITY_UNSUPPORTED" => Self::CapabilityUnsupported,
1098            "QUERY_OUTCOME_UNKNOWN" => Self::QueryOutcomeUnknown,
1099            "INVALID_QUERY_OPTIONS" => Self::InvalidQueryOptions,
1100            "INCOMPATIBLE_SQL_CONTROLS" => Self::IncompatibleSqlControls,
1101            "INVALID_IDEMPOTENCY_KEY" => Self::InvalidIdempotencyKey,
1102            "IDEMPOTENCY_KEY_REUSE_MISMATCH" => Self::IdempotencyKeyReuseMismatch,
1103            "IDEMPOTENCY_REQUIRES_JSON" => Self::IdempotencyRequiresJson,
1104            "IDEMPOTENCY_REQUIRES_SINGLE_WRITE" => Self::IdempotencyRequiresSingleWrite,
1105            "IDEMPOTENCY_STORE_FULL" => Self::IdempotencyStoreFull,
1106            "IDEMPOTENCY_UNSUPPORTED_IN_TRANSACTION" => Self::IdempotencyUnsupportedInTransaction,
1107            "IDEMPOTENCY_STORE_UNAVAILABLE" => Self::IdempotencyStoreUnavailable,
1108            "PAGINATION_REQUIRES_JSON" => Self::PaginationRequiresJson,
1109            "PAGINATION_REQUIRES_SINGLE_READ_QUERY" => Self::PaginationRequiresSingleReadQuery,
1110            "INVALID_PAGINATION_OPTIONS" => Self::InvalidPaginationOptions,
1111            "INVALID_SQL_CURSOR" => Self::InvalidSqlCursor,
1112            "SQL_CURSOR_EXPIRED" => Self::SqlCursorExpired,
1113            "SQL_CURSOR_NOT_FOUND" => Self::SqlCursorNotFound,
1114            "INVALID_SQL_PROJECTION" => Self::InvalidSqlProjection,
1115            "INVALID_PAGE_OFFSET" => Self::InvalidPageOffset,
1116            "SQL_PAGE_STORE_FULL" => Self::SqlPageStoreFull,
1117            "SQL_ADMISSION_CLOSED" => Self::SqlAdmissionClosed,
1118            "ENTROPY_UNAVAILABLE" => Self::EntropyUnavailable,
1119            other => Self::Other(other.into()),
1120        }
1121    }
1122}
1123
1124impl std::fmt::Display for RemoteQueryErrorCode {
1125    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1126        formatter.write_str(self.as_str())
1127    }
1128}
1129
1130impl<'de> Deserialize<'de> for RemoteQueryErrorCode {
1131    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
1132    where
1133        D: serde::Deserializer<'de>,
1134    {
1135        let code = String::deserialize(deserializer)?;
1136        Ok(Self::from_code(&code))
1137    }
1138}
1139
1140#[derive(Debug, Clone, Default, Deserialize)]
1141#[serde(deny_unknown_fields)]
1142pub struct RemoteQueryOutcome {
1143    #[serde(default)]
1144    pub committed: Option<bool>,
1145    #[serde(default)]
1146    pub committed_statements: Option<usize>,
1147    #[serde(default)]
1148    pub last_commit_epoch: Option<u64>,
1149    #[serde(default)]
1150    pub last_commit_epoch_text: Option<String>,
1151    #[serde(default)]
1152    pub first_commit_statement_index: Option<usize>,
1153    #[serde(default)]
1154    pub last_commit_statement_index: Option<usize>,
1155    #[serde(default)]
1156    pub completed_statements: Option<usize>,
1157    #[serde(default)]
1158    pub statement_index: Option<usize>,
1159    #[serde(default)]
1160    pub serialization: String,
1161}
1162
1163#[derive(Debug, Clone, Deserialize)]
1164#[serde(deny_unknown_fields)]
1165pub struct RemoteTerminalError {
1166    pub code: RemoteQueryErrorCode,
1167    pub category: String,
1168}
1169
1170#[derive(Debug, Clone, Deserialize)]
1171#[serde(deny_unknown_fields)]
1172pub struct RemoteQueryErrorBody {
1173    pub code: RemoteQueryErrorCode,
1174    pub message: String,
1175    #[serde(default)]
1176    pub query_id: Option<String>,
1177    #[serde(default)]
1178    pub committed: Option<bool>,
1179    #[serde(default)]
1180    pub retryable: bool,
1181}
1182
1183#[derive(Debug, Clone, Deserialize)]
1184#[serde(deny_unknown_fields)]
1185pub struct RemoteQueryErrorResponse {
1186    #[serde(default)]
1187    pub query_id: Option<String>,
1188    #[serde(default)]
1189    pub status: String,
1190    #[serde(default)]
1191    pub terminal_state: Option<String>,
1192    #[serde(default)]
1193    pub committed: Option<bool>,
1194    #[serde(default)]
1195    pub committed_statements: Option<usize>,
1196    #[serde(default)]
1197    pub last_commit_epoch: Option<u64>,
1198    #[serde(default)]
1199    pub last_commit_epoch_text: Option<String>,
1200    #[serde(default)]
1201    pub first_commit_statement_index: Option<usize>,
1202    #[serde(default)]
1203    pub last_commit_statement_index: Option<usize>,
1204    #[serde(default)]
1205    pub completed_statements: Option<usize>,
1206    #[serde(default)]
1207    pub statement_index: Option<usize>,
1208    #[serde(default)]
1209    pub cancel_outcome: Option<RemoteCancelOutcome>,
1210    #[serde(default)]
1211    pub cancellation_reason: Option<String>,
1212    #[serde(default)]
1213    pub retryable: bool,
1214    #[serde(default)]
1215    pub server_state: Option<String>,
1216    #[serde(default)]
1217    pub outcome: RemoteQueryOutcome,
1218    pub error: RemoteQueryErrorBody,
1219}
1220
1221#[derive(Debug, Clone, Deserialize)]
1222#[serde(deny_unknown_fields)]
1223pub struct SqlCancellationCapabilities {
1224    pub version: u8,
1225    pub client_query_ids: bool,
1226    pub cancel_endpoint: bool,
1227    pub query_status: bool,
1228    #[serde(default)]
1229    pub pre_registration_cancel: bool,
1230    pub stream_disconnect_cancels: bool,
1231}
1232
1233#[derive(Debug, Clone, Deserialize)]
1234#[serde(deny_unknown_fields)]
1235pub struct SqlIdempotencyCapabilities {
1236    pub version: u8,
1237    pub durable_pre_execution_intent: bool,
1238    pub replay_committed_receipt: bool,
1239    pub indeterminate_never_reexecutes: bool,
1240}
1241
1242#[derive(Debug, Clone, Deserialize)]
1243#[serde(deny_unknown_fields)]
1244pub struct SqlPaginationCapabilities {
1245    pub version: u8,
1246    pub continuation_endpoint: String,
1247    pub retained_snapshot: bool,
1248    pub projection_required: bool,
1249    pub byte_and_token_hints: bool,
1250}
1251
1252#[derive(Debug, Clone, Deserialize)]
1253#[serde(deny_unknown_fields)]
1254pub struct ServerCapabilities {
1255    pub sql_cancellation: SqlCancellationCapabilities,
1256    #[serde(default)]
1257    pub sql_idempotency: Option<SqlIdempotencyCapabilities>,
1258    #[serde(default)]
1259    pub sql_pagination: Option<SqlPaginationCapabilities>,
1260}
1261
1262#[derive(Debug, Clone, Deserialize)]
1263#[serde(deny_unknown_fields)]
1264pub struct RemoteQueryStatus {
1265    pub query_id: String,
1266    #[serde(default)]
1267    pub detail: Option<String>,
1268    #[serde(default)]
1269    pub status: String,
1270    #[serde(default)]
1271    pub state: String,
1272    #[serde(default)]
1273    pub server_state: String,
1274    #[serde(default)]
1275    pub terminal_state: Option<String>,
1276    #[serde(default)]
1277    pub operation: String,
1278    #[serde(default)]
1279    pub started_ms_ago: Option<u64>,
1280    #[serde(default)]
1281    pub deadline_ms_remaining: Option<u64>,
1282    #[serde(default)]
1283    pub session_id: Option<String>,
1284    #[serde(default)]
1285    pub code: Option<RemoteQueryErrorCode>,
1286    #[serde(default)]
1287    pub committed: Option<bool>,
1288    #[serde(default)]
1289    pub cancellation_reason: String,
1290    #[serde(default)]
1291    pub committed_statements: Option<usize>,
1292    #[serde(default)]
1293    pub last_commit_epoch: Option<u64>,
1294    #[serde(default)]
1295    pub last_commit_epoch_text: Option<String>,
1296    #[serde(default)]
1297    pub first_commit_statement_index: Option<usize>,
1298    #[serde(default)]
1299    pub last_commit_statement_index: Option<usize>,
1300    #[serde(default)]
1301    pub completed_statements: Option<usize>,
1302    #[serde(default)]
1303    pub statement_index: Option<usize>,
1304    #[serde(default)]
1305    pub cancel_outcome: Option<RemoteCancelOutcome>,
1306    #[serde(default)]
1307    pub retryable: bool,
1308    #[serde(default)]
1309    pub outcome: RemoteQueryOutcome,
1310    #[serde(default)]
1311    pub terminal_error: Option<RemoteTerminalError>,
1312    #[serde(default)]
1313    pub trace: serde_json::Value,
1314}
1315
1316impl RemoteQueryStatus {
1317    pub fn is_terminal(&self) -> bool {
1318        matches!(
1319            self.server_state_or_state(),
1320            "completed" | "failed" | "cancelled" | "pre_cancelled" | "finished"
1321        )
1322    }
1323
1324    pub fn durably_committed(&self) -> Option<bool> {
1325        match (self.committed, self.outcome.committed) {
1326            (Some(true), _) | (_, Some(true)) => Some(true),
1327            (Some(false), _) | (_, Some(false)) => Some(false),
1328            (None, None) => None,
1329        }
1330    }
1331
1332    pub fn server_state_or_state(&self) -> &str {
1333        if self.server_state.is_empty() {
1334            &self.state
1335        } else {
1336            &self.server_state
1337        }
1338    }
1339}
1340
1341#[derive(Debug, Clone, Deserialize)]
1342#[serde(deny_unknown_fields)]
1343pub struct RemoteSqlReceipt {
1344    pub query_id: String,
1345    pub original_query_id: String,
1346    pub status: String,
1347    #[serde(default)]
1348    pub terminal_state: Option<String>,
1349    #[serde(default)]
1350    pub server_state: String,
1351    #[serde(default)]
1352    pub cancel_outcome: Option<RemoteCancelOutcome>,
1353    #[serde(default)]
1354    pub cancellation_reason: String,
1355    pub committed: bool,
1356    pub committed_statements: usize,
1357    pub last_commit_epoch: Option<u64>,
1358    #[serde(default)]
1359    pub last_commit_epoch_text: Option<String>,
1360    #[serde(default)]
1361    pub first_commit_statement_index: Option<usize>,
1362    #[serde(default)]
1363    pub last_commit_statement_index: Option<usize>,
1364    #[serde(default)]
1365    pub completed_statements: usize,
1366    #[serde(default)]
1367    pub statement_index: usize,
1368    pub retryable: bool,
1369    pub idempotency_replayed: bool,
1370    pub idempotency_persisted: bool,
1371    pub idempotency_expires_at_ms: u64,
1372    pub outcome: RemoteQueryOutcome,
1373    #[serde(default)]
1374    pub terminal_error: Option<RemoteTerminalError>,
1375}
1376
1377fn exact_epoch(text: Option<&str>, numeric: Option<u64>) -> Result<Option<u64>, String> {
1378    match text {
1379        Some(text) => {
1380            let epoch = text
1381                .parse::<u64>()
1382                .map_err(|_| "last_commit_epoch_text is not an unsigned integer".to_owned())?;
1383            if epoch.to_string() != text {
1384                return Err("last_commit_epoch_text is not canonical".into());
1385            }
1386            if numeric.is_some_and(|numeric| numeric != epoch) {
1387                return Err("last_commit_epoch and last_commit_epoch_text disagree".into());
1388            }
1389            Ok(Some(epoch))
1390        }
1391        None => Ok(numeric),
1392    }
1393}
1394
1395fn validate_remote_query_status(
1396    mut status: RemoteQueryStatus,
1397    expected_query_id: mongreldb_query::QueryId,
1398) -> Result<RemoteQueryStatus, String> {
1399    const STATES: &[&str] = &[
1400        "queued",
1401        "planning",
1402        "executing",
1403        "streaming",
1404        "serializing",
1405        "commit_critical",
1406        "cancelling",
1407        "completed",
1408        "failed",
1409        "cancelled",
1410        "pre_cancelled",
1411        "finished",
1412    ];
1413    const STATUSES: &[&str] = &[
1414        "running",
1415        "outcome_unknown",
1416        "completed",
1417        "failed_before_commit",
1418        "cancelled_before_commit",
1419        "deadline_before_commit",
1420        "cancelled_before_start",
1421        "committed",
1422        "committed_with_error",
1423        "partially_committed",
1424        "cancelled_after_commit",
1425        "deadline_after_commit",
1426        "finished",
1427    ];
1428    if status.query_id != expected_query_id.to_string() {
1429        return Err("query status query_id does not match the request".into());
1430    }
1431    if status
1432        .detail
1433        .as_deref()
1434        .is_some_and(|detail| detail != "compact")
1435    {
1436        return Err("query status detail is invalid".into());
1437    }
1438    if !STATUSES.contains(&status.status.as_str())
1439        || status.state.is_empty()
1440        || !STATES.contains(&status.state.as_str())
1441        || (!status.server_state.is_empty()
1442            && (!STATES.contains(&status.server_state.as_str())
1443                || status.server_state != status.state))
1444        || status
1445            .terminal_state
1446            .as_ref()
1447            .is_some_and(|terminal| terminal != &status.status)
1448    {
1449        return Err("query status state or status is invalid".into());
1450    }
1451    let state_matches_status = match status.status.as_str() {
1452        "running" => matches!(
1453            status.state.as_str(),
1454            "queued"
1455                | "planning"
1456                | "executing"
1457                | "streaming"
1458                | "serializing"
1459                | "commit_critical"
1460                | "cancelling"
1461        ),
1462        "committed" => !matches!(
1463            status.state.as_str(),
1464            "failed" | "cancelled" | "pre_cancelled" | "finished"
1465        ),
1466        "completed" => status.state == "completed",
1467        "failed_before_commit"
1468        | "committed_with_error"
1469        | "partially_committed"
1470        | "outcome_unknown" => status.state == "failed",
1471        "cancelled_before_commit"
1472        | "deadline_before_commit"
1473        | "cancelled_after_commit"
1474        | "deadline_after_commit" => status.state == "cancelled",
1475        "cancelled_before_start" => status.state == "pre_cancelled",
1476        "finished" => status.state == "finished",
1477        _ => false,
1478    };
1479    if !state_matches_status {
1480        return Err("query status state and status disagree".into());
1481    }
1482    let top_epoch = exact_epoch(
1483        status.last_commit_epoch_text.as_deref(),
1484        status.last_commit_epoch,
1485    )?;
1486    let outcome_epoch = exact_epoch(
1487        status.outcome.last_commit_epoch_text.as_deref(),
1488        status.outcome.last_commit_epoch,
1489    )?;
1490    if status
1491        .last_commit_epoch
1492        .is_some_and(|numeric| Some(numeric) != top_epoch)
1493        || status
1494            .outcome
1495            .last_commit_epoch
1496            .is_some_and(|numeric| Some(numeric) != outcome_epoch)
1497        || top_epoch != outcome_epoch
1498        || status.committed != status.outcome.committed
1499        || status.committed_statements != status.outcome.committed_statements
1500        || status.first_commit_statement_index != status.outcome.first_commit_statement_index
1501        || status.last_commit_statement_index != status.outcome.last_commit_statement_index
1502        || status.completed_statements != status.outcome.completed_statements
1503        || status.statement_index != status.outcome.statement_index
1504    {
1505        return Err("query status top-level and outcome fields disagree".into());
1506    }
1507    match status.committed {
1508        Some(true) => {
1509            if status.committed_statements == Some(0)
1510                || status.committed_statements.is_none()
1511                || top_epoch.is_none()
1512                || status.last_commit_epoch_text.is_none()
1513                || status.outcome.last_commit_epoch_text.is_none()
1514                || status.first_commit_statement_index.is_none()
1515                || status.last_commit_statement_index.is_none()
1516                || status.completed_statements.is_none()
1517                || status.statement_index.is_none()
1518                || !matches!(
1519                    status.status.as_str(),
1520                    "committed"
1521                        | "committed_with_error"
1522                        | "partially_committed"
1523                        | "cancelled_after_commit"
1524                        | "deadline_after_commit"
1525                )
1526            {
1527                return Err("committed query status has invalid durable metadata".into());
1528            }
1529        }
1530        Some(false) => {
1531            if status.committed_statements != Some(0)
1532                || top_epoch.is_some()
1533                || status.first_commit_statement_index.is_some()
1534                || status.last_commit_statement_index.is_some()
1535                || status.completed_statements.is_none()
1536                || status.statement_index.is_none()
1537                || matches!(
1538                    status.status.as_str(),
1539                    "committed"
1540                        | "committed_with_error"
1541                        | "partially_committed"
1542                        | "cancelled_after_commit"
1543                        | "deadline_after_commit"
1544                        | "outcome_unknown"
1545                        | "finished"
1546                )
1547            {
1548                return Err("non-committed query status has invalid durable metadata".into());
1549            }
1550        }
1551        None => {
1552            if status.committed_statements.is_some()
1553                || top_epoch.is_some()
1554                || status.first_commit_statement_index.is_some()
1555                || status.last_commit_statement_index.is_some()
1556                || status.completed_statements.is_some()
1557                || status.statement_index.is_some()
1558                || !matches!(status.status.as_str(), "outcome_unknown" | "finished")
1559            {
1560                return Err("unknown query status contains durable metadata".into());
1561            }
1562        }
1563    }
1564    if let (Some(first), Some(last), Some(committed)) = (
1565        status.first_commit_statement_index,
1566        status.last_commit_statement_index,
1567        status.committed_statements,
1568    ) {
1569        if first > last
1570            || committed > last.saturating_sub(first).saturating_add(1)
1571            || status
1572                .statement_index
1573                .is_some_and(|statement| last > statement)
1574        {
1575            return Err("query status commit statement indexes are invalid".into());
1576        }
1577    }
1578    if let (Some(completed), Some(statement)) =
1579        (status.completed_statements, status.statement_index)
1580    {
1581        if statement > completed || completed > statement.saturating_add(1) {
1582            return Err("query status statement index and completed count disagree".into());
1583        }
1584    }
1585    let terminal_error = status.terminal_error.as_ref();
1586    if terminal_error.is_some_and(|error| {
1587        error.code.as_str().trim().is_empty()
1588            || !matches!(
1589                error.category.as_str(),
1590                "cancellation" | "deadline" | "result_limit" | "serialization" | "execution"
1591            )
1592    }) {
1593        return Err("query status terminal error fields are invalid".into());
1594    }
1595    let terminal_error_matches = match status.status.as_str() {
1596        "running" | "completed" | "committed" | "finished" => terminal_error.is_none(),
1597        "outcome_unknown" => terminal_error.is_some_and(|error| {
1598            error.code.as_str() == "QUERY_OUTCOME_UNKNOWN" && error.category == "execution"
1599        }),
1600        "cancelled_before_start" | "cancelled_before_commit" => {
1601            terminal_error.is_some_and(|error| {
1602                error.code.as_str() == "QUERY_CANCELLED" && error.category == "cancellation"
1603            })
1604        }
1605        "cancelled_after_commit" => terminal_error.is_some_and(|error| {
1606            error.code.as_str() == "QUERY_CANCELLED_AFTER_COMMIT"
1607                && error.category == "cancellation"
1608        }),
1609        "deadline_before_commit" => terminal_error.is_some_and(|error| {
1610            error.code.as_str() == "DEADLINE_EXCEEDED" && error.category == "deadline"
1611        }),
1612        "deadline_after_commit" => terminal_error.is_some_and(|error| {
1613            error.code.as_str() == "DEADLINE_AFTER_COMMIT" && error.category == "deadline"
1614        }),
1615        "failed_before_commit" | "committed_with_error" | "partially_committed" => {
1616            terminal_error.is_some()
1617        }
1618        _ => false,
1619    };
1620    if !terminal_error_matches {
1621        return Err("query status terminal error disagrees with status".into());
1622    }
1623    if terminal_error.is_some_and(|error| {
1624        (error.category == "cancellation")
1625            != matches!(
1626                error.code.as_str(),
1627                "QUERY_CANCELLED" | "QUERY_CANCELLED_AFTER_COMMIT"
1628            )
1629            || (error.category == "deadline")
1630                != matches!(
1631                    error.code.as_str(),
1632                    "DEADLINE_EXCEEDED" | "DEADLINE_AFTER_COMMIT"
1633                )
1634    }) {
1635        return Err("query status terminal error code and category disagree".into());
1636    }
1637    let retryable = terminal_error.is_some_and(|error| {
1638        matches!(
1639            error.code.as_str(),
1640            "IDEMPOTENCY_STORE_FULL" | "IDEMPOTENCY_STORE_UNAVAILABLE"
1641        )
1642    });
1643    if status.retryable != retryable {
1644        return Err("query status retryable flag disagrees with terminal error".into());
1645    }
1646    let expected_cancel_outcome = match status.state.as_str() {
1647        "commit_critical" => Some(RemoteCancelOutcome::TooLate),
1648        "cancelling" => Some(RemoteCancelOutcome::Accepted),
1649        "completed" | "failed" | "cancelled" | "finished" => {
1650            Some(RemoteCancelOutcome::AlreadyFinished)
1651        }
1652        "pre_cancelled" => Some(RemoteCancelOutcome::PreCancelled),
1653        _ => None,
1654    };
1655    if status.cancel_outcome != expected_cancel_outcome {
1656        return Err("query status cancel_outcome disagrees with state".into());
1657    }
1658    const CANCELLATION_REASONS: &[&str] = &[
1659        "none",
1660        "client_request",
1661        "client_disconnected",
1662        "session_closed",
1663        "server_shutdown",
1664        "deadline",
1665    ];
1666    let valid_reason = match status.status.as_str() {
1667        "finished" => status.cancellation_reason == "none",
1668        "cancelled_before_start" | "cancelled_before_commit" | "cancelled_after_commit" => {
1669            CANCELLATION_REASONS.contains(&status.cancellation_reason.as_str())
1670                && !matches!(status.cancellation_reason.as_str(), "none" | "deadline")
1671        }
1672        "deadline_before_commit" | "deadline_after_commit" => {
1673            status.cancellation_reason == "deadline"
1674        }
1675        "running" | "committed" if status.state == "cancelling" => {
1676            CANCELLATION_REASONS.contains(&status.cancellation_reason.as_str())
1677                && status.cancellation_reason != "none"
1678        }
1679        _ => status.cancellation_reason == "none",
1680    };
1681    if !valid_reason {
1682        return Err("query status cancellation_reason disagrees with status".into());
1683    }
1684    let serialization = status.outcome.serialization.as_str();
1685    let valid_serialization = match status.status.as_str() {
1686        "finished" | "outcome_unknown" => serialization == "unknown",
1687        "completed" => serialization == "succeeded",
1688        "running" | "committed" => match status.state.as_str() {
1689            "serializing" => serialization == "in_progress",
1690            "cancelling" => matches!(serialization, "not_started" | "in_progress"),
1691            "completed" => serialization == "succeeded",
1692            _ => serialization == "not_started",
1693        },
1694        _ => matches!(serialization, "not_started" | "failed"),
1695    };
1696    if !valid_serialization {
1697        return Err("query status serialization state is invalid".into());
1698    }
1699    let terminal_state_valid = match status.status.as_str() {
1700        "running" | "finished" => status.terminal_state.is_none(),
1701        "committed" if status.state != "completed" => status.terminal_state.is_none(),
1702        _ => status.terminal_state.as_deref() == Some(status.status.as_str()),
1703    };
1704    if !terminal_state_valid {
1705        return Err("query status terminal_state is invalid".into());
1706    }
1707    if matches!(status.state.as_str(), "pre_cancelled" | "finished") {
1708        if !status.operation.is_empty() {
1709            return Err("synthetic query status unexpectedly names an operation".into());
1710        }
1711    } else if status.operation.is_empty() {
1712        return Err("live query status lacks operation".into());
1713    }
1714    status.last_commit_epoch = top_epoch;
1715    status.outcome.last_commit_epoch = outcome_epoch;
1716    Ok(status)
1717}
1718
1719fn validate_remote_query_error(
1720    mut response: RemoteQueryErrorResponse,
1721    expected_query_id: mongreldb_query::QueryId,
1722) -> Result<RemoteQueryErrorResponse, String> {
1723    let expected = expected_query_id.to_string();
1724    if response.query_id.as_deref() != Some(&expected)
1725        || response.error.query_id.as_deref() != Some(&expected)
1726        || response.terminal_state.as_deref() != Some(response.status.as_str())
1727        || response.error.message.is_empty()
1728    {
1729        return Err("query error query_id does not match the request".into());
1730    }
1731    if response.committed != response.outcome.committed
1732        || response.committed != response.error.committed
1733        || response.committed_statements != response.outcome.committed_statements
1734        || response.first_commit_statement_index != response.outcome.first_commit_statement_index
1735        || response.last_commit_statement_index != response.outcome.last_commit_statement_index
1736        || response.completed_statements != response.outcome.completed_statements
1737        || response.statement_index != response.outcome.statement_index
1738        || response.retryable != response.error.retryable
1739    {
1740        return Err("query error top-level, outcome, and error fields disagree".into());
1741    }
1742    let top_epoch = exact_epoch(
1743        response.last_commit_epoch_text.as_deref(),
1744        response.last_commit_epoch,
1745    )?;
1746    let outcome_epoch = exact_epoch(
1747        response.outcome.last_commit_epoch_text.as_deref(),
1748        response.outcome.last_commit_epoch,
1749    )?;
1750    if response
1751        .last_commit_epoch
1752        .is_some_and(|numeric| Some(numeric) != top_epoch)
1753        || response
1754            .outcome
1755            .last_commit_epoch
1756            .is_some_and(|numeric| Some(numeric) != outcome_epoch)
1757        || top_epoch != outcome_epoch
1758    {
1759        return Err("query error top-level and outcome commit epochs disagree".into());
1760    }
1761    let outcome_unknown = response.error.code == RemoteQueryErrorCode::QueryOutcomeUnknown;
1762    match response.committed {
1763        Some(true) => {
1764            if outcome_unknown
1765                || response.committed_statements == Some(0)
1766                || response.committed_statements.is_none()
1767                || top_epoch.is_none()
1768                || response.last_commit_epoch_text.is_none()
1769                || response.outcome.last_commit_epoch_text.is_none()
1770                || response.first_commit_statement_index.is_none()
1771                || response.last_commit_statement_index.is_none()
1772                || response.completed_statements.is_none()
1773                || response.statement_index.is_none()
1774                || !matches!(
1775                    response.status.as_str(),
1776                    "committed"
1777                        | "committed_with_error"
1778                        | "partially_committed"
1779                        | "cancelled_after_commit"
1780                        | "deadline_after_commit"
1781                )
1782            {
1783                return Err("committed query error has invalid durable metadata".into());
1784            }
1785        }
1786        Some(false) => {
1787            if outcome_unknown
1788                || response.committed_statements != Some(0)
1789                || top_epoch.is_some()
1790                || response.first_commit_statement_index.is_some()
1791                || response.last_commit_statement_index.is_some()
1792                || response.completed_statements.is_none()
1793                || response.statement_index.is_none()
1794                || !matches!(
1795                    response.status.as_str(),
1796                    "failed_before_commit"
1797                        | "cancelled_before_commit"
1798                        | "deadline_before_commit"
1799                        | "cancelled_before_start"
1800                )
1801            {
1802                return Err("non-committed query error has invalid durable metadata".into());
1803            }
1804        }
1805        None => {
1806            if !outcome_unknown
1807                || response.status != "outcome_unknown"
1808                || response.committed_statements.is_some()
1809                || top_epoch.is_some()
1810                || response.first_commit_statement_index.is_some()
1811                || response.last_commit_statement_index.is_some()
1812                || response.completed_statements.is_some()
1813                || response.statement_index.is_some()
1814                || response.retryable
1815            {
1816                return Err("unknown query error contains contradictory metadata".into());
1817            }
1818        }
1819    }
1820    if outcome_unknown && response.status != "outcome_unknown" {
1821        return Err("outcome-unknown error has the wrong status".into());
1822    }
1823    if response.retryable
1824        && (response.committed != Some(false)
1825            || !matches!(
1826                response.error.code.as_str(),
1827                "QUERY_REGISTRY_FULL" | "IDEMPOTENCY_STORE_FULL" | "IDEMPOTENCY_STORE_UNAVAILABLE"
1828            ))
1829    {
1830        return Err("query error retryable flag is unsafe".into());
1831    }
1832    if let (Some(first), Some(last), Some(committed)) = (
1833        response.first_commit_statement_index,
1834        response.last_commit_statement_index,
1835        response.committed_statements,
1836    ) {
1837        if first > last
1838            || committed > last.saturating_sub(first).saturating_add(1)
1839            || response
1840                .statement_index
1841                .is_some_and(|statement| last > statement)
1842        {
1843            return Err("query error commit statement indexes are invalid".into());
1844        }
1845    }
1846    if let (Some(completed), Some(statement)) =
1847        (response.completed_statements, response.statement_index)
1848    {
1849        if statement > completed || completed > statement.saturating_add(1) {
1850            return Err("query error statement index and completed count disagree".into());
1851        }
1852    }
1853    let code_matches = match response.error.code {
1854        RemoteQueryErrorCode::QueryOutcomeUnknown => response.status == "outcome_unknown",
1855        RemoteQueryErrorCode::QueryCancelledAfterCommit => {
1856            response.status == "cancelled_after_commit" && response.committed == Some(true)
1857        }
1858        RemoteQueryErrorCode::DeadlineAfterCommit => {
1859            response.status == "deadline_after_commit" && response.committed == Some(true)
1860        }
1861        RemoteQueryErrorCode::QueryCancelled => matches!(
1862            response.status.as_str(),
1863            "cancelled_before_commit" | "cancelled_before_start"
1864        ),
1865        RemoteQueryErrorCode::DeadlineExceeded => response.status == "deadline_before_commit",
1866        RemoteQueryErrorCode::CommitOutcome
1867        | RemoteQueryErrorCode::SerializationFailedAfterCommit => response.committed == Some(true),
1868        RemoteQueryErrorCode::SerializationFailed => response.committed == Some(false),
1869        _ => true,
1870    };
1871    if !code_matches {
1872        return Err("query error code and status disagree".into());
1873    }
1874    let status_matches_code = match response.status.as_str() {
1875        "outcome_unknown" => response.error.code == RemoteQueryErrorCode::QueryOutcomeUnknown,
1876        "cancelled_after_commit" => {
1877            response.error.code == RemoteQueryErrorCode::QueryCancelledAfterCommit
1878        }
1879        "deadline_after_commit" => response.error.code == RemoteQueryErrorCode::DeadlineAfterCommit,
1880        "cancelled_before_commit" | "cancelled_before_start" => {
1881            response.error.code == RemoteQueryErrorCode::QueryCancelled
1882        }
1883        "deadline_before_commit" => response.error.code == RemoteQueryErrorCode::DeadlineExceeded,
1884        _ => true,
1885    };
1886    if !status_matches_code {
1887        return Err("query error status and code disagree".into());
1888    }
1889    const CANCELLATION_REASONS: &[&str] = &[
1890        "none",
1891        "client_request",
1892        "client_disconnected",
1893        "session_closed",
1894        "server_shutdown",
1895        "deadline",
1896    ];
1897    let cancellation_error = matches!(
1898        response.error.code,
1899        RemoteQueryErrorCode::QueryCancelled | RemoteQueryErrorCode::QueryCancelledAfterCommit
1900    );
1901    let deadline_error = matches!(
1902        response.error.code,
1903        RemoteQueryErrorCode::DeadlineExceeded | RemoteQueryErrorCode::DeadlineAfterCommit
1904    );
1905    let expected_cancel_outcome = if cancellation_error || deadline_error {
1906        Some(RemoteCancelOutcome::Accepted)
1907    } else {
1908        match response.server_state.as_deref() {
1909            Some("commit_critical") => Some(RemoteCancelOutcome::TooLate),
1910            Some("cancelling") => Some(RemoteCancelOutcome::Accepted),
1911            Some("completed" | "failed" | "cancelled") => {
1912                Some(RemoteCancelOutcome::AlreadyFinished)
1913            }
1914            Some(_) | None => None,
1915        }
1916    };
1917    if response.cancel_outcome != expected_cancel_outcome {
1918        return Err("query error cancel_outcome disagrees with its phase".into());
1919    }
1920    let valid_reason = match response.cancellation_reason.as_deref() {
1921        Some("deadline") => deadline_error,
1922        Some(reason) if CANCELLATION_REASONS.contains(&reason) => {
1923            if cancellation_error {
1924                reason != "none"
1925            } else {
1926                response.server_state.is_some() && reason == "none"
1927            }
1928        }
1929        None => response.server_state.is_none() && !cancellation_error && !deadline_error,
1930        Some(_) => false,
1931    };
1932    if !valid_reason {
1933        return Err("query error cancellation_reason is invalid".into());
1934    }
1935    let valid_server_state = match response.server_state.as_deref() {
1936        None => true,
1937        Some("cancelled") => cancellation_error || deadline_error,
1938        Some("failed") => !cancellation_error && !deadline_error,
1939        Some(_) => false,
1940    };
1941    if !valid_server_state {
1942        return Err("query error server_state disagrees with its status".into());
1943    }
1944    let valid_serialization = match response.server_state.as_deref() {
1945        None => response.outcome.serialization == "unknown",
1946        Some(_) if outcome_unknown => response.outcome.serialization == "unknown",
1947        Some("failed" | "cancelled") => {
1948            matches!(
1949                response.outcome.serialization.as_str(),
1950                "not_started" | "failed"
1951            )
1952        }
1953        Some(_) => false,
1954    };
1955    if !valid_serialization {
1956        return Err("query error serialization state is invalid".into());
1957    }
1958    response.last_commit_epoch = top_epoch;
1959    response.outcome.last_commit_epoch = outcome_epoch;
1960    Ok(response)
1961}
1962
1963fn validate_queryless_sql_error(
1964    mut response: RemoteQueryErrorResponse,
1965) -> Result<RemoteQueryErrorResponse, String> {
1966    if response.query_id.is_some()
1967        || response.error.query_id.is_some()
1968        || response.status != "failed_before_commit"
1969        || response.terminal_state.as_deref() != Some("failed_before_commit")
1970        || response.server_state.as_deref() != Some("failed")
1971        || response.committed != Some(false)
1972        || response.error.committed != Some(false)
1973        || response.outcome.committed != Some(false)
1974        || response.committed_statements != Some(0)
1975        || response.outcome.committed_statements != Some(0)
1976        || response.last_commit_epoch.is_some()
1977        || response.last_commit_epoch_text.is_some()
1978        || response.outcome.last_commit_epoch.is_some()
1979        || response.outcome.last_commit_epoch_text.is_some()
1980        || response.first_commit_statement_index.is_some()
1981        || response.last_commit_statement_index.is_some()
1982        || response.outcome.first_commit_statement_index.is_some()
1983        || response.outcome.last_commit_statement_index.is_some()
1984        || response.completed_statements != Some(0)
1985        || response.outcome.completed_statements != Some(0)
1986        || response.statement_index != Some(0)
1987        || response.outcome.statement_index != Some(0)
1988        || response.cancel_outcome.is_some()
1989        || response.cancellation_reason.is_some()
1990        || response.retryable
1991        || response.error.retryable
1992        || response.outcome.serialization != "not_started"
1993        || response.error.message.is_empty()
1994        || !matches!(
1995            response.error.code,
1996            RemoteQueryErrorCode::InvalidSqlCursor
1997                | RemoteQueryErrorCode::SqlCursorExpired
1998                | RemoteQueryErrorCode::SqlCursorNotFound
1999                | RemoteQueryErrorCode::ResultLimitExceeded
2000                | RemoteQueryErrorCode::SerializationFailed
2001                | RemoteQueryErrorCode::SerializationWorkerFailed
2002                | RemoteQueryErrorCode::SqlAdmissionClosed
2003                | RemoteQueryErrorCode::EntropyUnavailable
2004        )
2005    {
2006        return Err("queryless SQL error metadata is invalid".into());
2007    }
2008    response.last_commit_epoch = None;
2009    response.outcome.last_commit_epoch = None;
2010    Ok(response)
2011}
2012
2013fn validate_remote_sql_receipt(
2014    mut receipt: RemoteSqlReceipt,
2015    expected_query_id: mongreldb_query::QueryId,
2016    expected_original_query_id: Option<mongreldb_query::QueryId>,
2017) -> Result<RemoteSqlReceipt, String> {
2018    if receipt.query_id != expected_query_id.to_string()
2019        || receipt.terminal_state.as_deref() != Some(receipt.status.as_str())
2020    {
2021        return Err("receipt query_id does not match the request".into());
2022    }
2023    if receipt
2024        .original_query_id
2025        .parse::<mongreldb_query::QueryId>()
2026        .is_err()
2027    {
2028        return Err("receipt original_query_id is invalid".into());
2029    }
2030    let status_committed = match receipt.status.as_str() {
2031        "completed" => false,
2032        "committed"
2033        | "committed_with_error"
2034        | "partially_committed"
2035        | "cancelled_after_commit"
2036        | "deadline_after_commit" => true,
2037        _ => return Err("receipt status is invalid".into()),
2038    };
2039    let expected_server_state = match receipt.status.as_str() {
2040        "completed" | "committed" => "completed",
2041        "committed_with_error" | "partially_committed" => "failed",
2042        "cancelled_after_commit" | "deadline_after_commit" => "cancelled",
2043        _ => return Err("receipt status is invalid".into()),
2044    };
2045    if receipt.server_state != expected_server_state
2046        || receipt.cancel_outcome != Some(RemoteCancelOutcome::AlreadyFinished)
2047    {
2048        return Err("receipt terminal state metadata is invalid".into());
2049    }
2050    let terminal_error = receipt.terminal_error.as_ref();
2051    let terminal_error_matches = match receipt.status.as_str() {
2052        "completed" | "committed" => terminal_error.is_none(),
2053        "cancelled_after_commit" => terminal_error.is_some_and(|error| {
2054            error.code.as_str() == "QUERY_CANCELLED_AFTER_COMMIT"
2055                && error.category == "cancellation"
2056        }),
2057        "deadline_after_commit" => terminal_error.is_some_and(|error| {
2058            error.code.as_str() == "DEADLINE_AFTER_COMMIT" && error.category == "deadline"
2059        }),
2060        "committed_with_error" | "partially_committed" => terminal_error.is_some(),
2061        _ => false,
2062    };
2063    if !terminal_error_matches
2064        || terminal_error.is_some_and(|error| {
2065            !matches!(
2066                error.category.as_str(),
2067                "cancellation" | "deadline" | "result_limit" | "serialization" | "execution"
2068            ) || (error.category == "cancellation")
2069                != (error.code.as_str() == "QUERY_CANCELLED_AFTER_COMMIT")
2070                || (error.category == "deadline")
2071                    != (error.code.as_str() == "DEADLINE_AFTER_COMMIT")
2072                || (error.category == "serialization")
2073                    != (error.code.as_str() == "SERIALIZATION_FAILED_AFTER_COMMIT")
2074        })
2075    {
2076        return Err("receipt terminal error disagrees with status".into());
2077    }
2078    let valid_reason = match receipt.status.as_str() {
2079        "cancelled_after_commit" => matches!(
2080            receipt.cancellation_reason.as_str(),
2081            "client_request" | "client_disconnected" | "session_closed" | "server_shutdown"
2082        ),
2083        "deadline_after_commit" => receipt.cancellation_reason == "deadline",
2084        _ => receipt.cancellation_reason == "none",
2085    };
2086    let serialization = receipt.outcome.serialization.as_str();
2087    let valid_serialization = matches!(serialization, "not_started" | "succeeded" | "failed")
2088        && match receipt.status.as_str() {
2089            "completed" | "committed" => serialization == "succeeded",
2090            _ if terminal_error.is_some_and(|error| error.category == "serialization") => {
2091                serialization == "failed"
2092            }
2093            _ => serialization != "succeeded",
2094        };
2095    if !valid_reason || !valid_serialization {
2096        return Err("receipt cancellation or serialization metadata is invalid".into());
2097    }
2098    if receipt.committed != status_committed
2099        || receipt.outcome.committed != Some(receipt.committed)
2100        || receipt.outcome.committed_statements != Some(receipt.committed_statements)
2101        || receipt.outcome.first_commit_statement_index != receipt.first_commit_statement_index
2102        || receipt.outcome.last_commit_statement_index != receipt.last_commit_statement_index
2103        || receipt.outcome.completed_statements != Some(receipt.completed_statements)
2104        || receipt.outcome.statement_index != Some(receipt.statement_index)
2105    {
2106        return Err("receipt top-level and outcome fields disagree".into());
2107    }
2108    let top_epoch = exact_epoch(
2109        receipt.last_commit_epoch_text.as_deref(),
2110        receipt.last_commit_epoch,
2111    )?;
2112    let outcome_epoch = exact_epoch(
2113        receipt.outcome.last_commit_epoch_text.as_deref(),
2114        receipt.outcome.last_commit_epoch,
2115    )?;
2116    if receipt
2117        .last_commit_epoch
2118        .is_some_and(|numeric| Some(numeric) != top_epoch)
2119        || receipt
2120            .outcome
2121            .last_commit_epoch
2122            .is_some_and(|numeric| Some(numeric) != outcome_epoch)
2123        || top_epoch != outcome_epoch
2124    {
2125        return Err("receipt top-level and outcome commit epochs disagree".into());
2126    }
2127    if receipt.committed {
2128        if receipt.committed_statements == 0
2129            || top_epoch.is_none()
2130            || receipt.last_commit_epoch_text.is_none()
2131            || receipt.outcome.last_commit_epoch_text.is_none()
2132            || receipt.first_commit_statement_index.is_none()
2133            || receipt.last_commit_statement_index.is_none()
2134        {
2135            return Err("committed receipt has no durable commit metadata".into());
2136        }
2137    } else if receipt.committed_statements != 0
2138        || top_epoch.is_some()
2139        || receipt.first_commit_statement_index.is_some()
2140        || receipt.last_commit_statement_index.is_some()
2141    {
2142        return Err("non-committed receipt contains commit metadata".into());
2143    }
2144    if let (Some(first), Some(last)) = (
2145        receipt.first_commit_statement_index,
2146        receipt.last_commit_statement_index,
2147    ) {
2148        if first > last
2149            || receipt.committed_statements > last.saturating_sub(first).saturating_add(1)
2150            || last > receipt.statement_index
2151        {
2152            return Err("receipt commit statement indexes are invalid".into());
2153        }
2154    }
2155    if receipt.statement_index > receipt.completed_statements
2156        || receipt.completed_statements > receipt.statement_index.saturating_add(1)
2157    {
2158        return Err("receipt statement index and completed count disagree".into());
2159    }
2160    let idempotency_identity_valid = match expected_original_query_id {
2161        Some(original) if receipt.idempotency_replayed => {
2162            receipt.original_query_id == original.to_string()
2163        }
2164        Some(_) => receipt.original_query_id == expected_query_id.to_string(),
2165        None if receipt.idempotency_replayed => true,
2166        None => receipt.original_query_id == expected_query_id.to_string(),
2167    };
2168    if !receipt.idempotency_persisted
2169        || receipt.idempotency_expires_at_ms == 0
2170        || receipt.retryable
2171        || !idempotency_identity_valid
2172    {
2173        return Err("receipt idempotency metadata is invalid".into());
2174    }
2175    receipt.last_commit_epoch = top_epoch;
2176    receipt.outcome.last_commit_epoch = outcome_epoch;
2177    Ok(receipt)
2178}
2179
2180#[derive(Clone)]
2181struct SqlReceiptCommitProof {
2182    epoch: u64,
2183    epoch_text: String,
2184    committed_statements: usize,
2185}
2186
2187fn sql_receipt_commit_proof(
2188    value: &serde_json::Value,
2189    expected_query_id: mongreldb_query::QueryId,
2190    expected_original_query_id: Option<mongreldb_query::QueryId>,
2191) -> Option<SqlReceiptCommitProof> {
2192    let object = value.as_object()?;
2193    let status = object.get("status")?.as_str()?;
2194    if !matches!(
2195        status,
2196        "committed"
2197            | "committed_with_error"
2198            | "partially_committed"
2199            | "cancelled_after_commit"
2200            | "deadline_after_commit"
2201    ) || object.get("query_id")?.as_str()? != expected_query_id.to_string()
2202        || !object.get("committed")?.as_bool()?
2203        || object.get("retryable")?.as_bool()?
2204        || !object.get("idempotency_persisted")?.as_bool()?
2205    {
2206        return None;
2207    }
2208    let original = object.get("original_query_id")?.as_str()?;
2209    let replayed = object.get("idempotency_replayed")?.as_bool()?;
2210    let identity_valid = match expected_original_query_id {
2211        Some(expected) if replayed => original == expected.to_string(),
2212        Some(_) => original == expected_query_id.to_string(),
2213        None if replayed => original.parse::<mongreldb_query::QueryId>().is_ok(),
2214        None => original == expected_query_id.to_string(),
2215    };
2216    if !identity_valid {
2217        return None;
2218    }
2219    let committed_statements =
2220        usize::try_from(object.get("committed_statements")?.as_u64()?).ok()?;
2221    if committed_statements == 0 {
2222        return None;
2223    }
2224    let epoch_text = object.get("last_commit_epoch_text")?.as_str()?.to_owned();
2225    let numeric_epoch = object.get("last_commit_epoch")?.as_u64();
2226    let epoch = exact_epoch(Some(&epoch_text), numeric_epoch).ok()??;
2227    if numeric_epoch.is_some_and(|numeric| numeric != epoch) {
2228        return None;
2229    }
2230    let outcome = object.get("outcome")?.as_object()?;
2231    let outcome_numeric_epoch = outcome
2232        .get("last_commit_epoch")
2233        .and_then(serde_json::Value::as_u64);
2234    let outcome_epoch = exact_epoch(
2235        outcome
2236            .get("last_commit_epoch_text")
2237            .and_then(serde_json::Value::as_str),
2238        outcome_numeric_epoch,
2239    )
2240    .ok()??;
2241    if !outcome.get("committed")?.as_bool()?
2242        || usize::try_from(outcome.get("committed_statements")?.as_u64()?).ok()?
2243            != committed_statements
2244        || outcome.get("last_commit_epoch_text")?.as_str()? != epoch_text
2245        || outcome_epoch != epoch
2246        || outcome_numeric_epoch.is_some_and(|numeric| numeric != outcome_epoch)
2247    {
2248        return None;
2249    }
2250    Some(SqlReceiptCommitProof {
2251        epoch,
2252        epoch_text,
2253        committed_statements,
2254    })
2255}
2256
2257fn committed_sql_receipt_decode_error(
2258    query_id: mongreldb_query::QueryId,
2259    proof: SqlReceiptCommitProof,
2260    message: impl Into<String>,
2261) -> ClientError {
2262    let message = message.into();
2263    let query_id = query_id.to_string();
2264    let code = RemoteQueryErrorCode::CommitOutcome;
2265    let response = RemoteQueryErrorResponse {
2266        query_id: Some(query_id.clone()),
2267        status: "committed_with_error".into(),
2268        terminal_state: Some("committed_with_error".into()),
2269        committed: Some(true),
2270        committed_statements: Some(proof.committed_statements),
2271        last_commit_epoch: Some(proof.epoch),
2272        last_commit_epoch_text: Some(proof.epoch_text.clone()),
2273        first_commit_statement_index: None,
2274        last_commit_statement_index: None,
2275        completed_statements: None,
2276        statement_index: None,
2277        cancel_outcome: Some(RemoteCancelOutcome::AlreadyFinished),
2278        cancellation_reason: Some("none".into()),
2279        retryable: false,
2280        server_state: Some("failed".into()),
2281        outcome: RemoteQueryOutcome {
2282            committed: Some(true),
2283            committed_statements: Some(proof.committed_statements),
2284            last_commit_epoch: Some(proof.epoch),
2285            last_commit_epoch_text: Some(proof.epoch_text),
2286            serialization: "unknown".into(),
2287            ..RemoteQueryOutcome::default()
2288        },
2289        error: RemoteQueryErrorBody {
2290            code: code.clone(),
2291            message: message.clone(),
2292            query_id: Some(query_id),
2293            committed: Some(true),
2294            retryable: false,
2295        },
2296    };
2297    ClientError::Query {
2298        status: 0,
2299        code,
2300        message,
2301        response: Box::new(response),
2302    }
2303}
2304
2305fn sql_error_commit_proof(
2306    value: &serde_json::Value,
2307    expected_query_id: mongreldb_query::QueryId,
2308) -> Option<SqlReceiptCommitProof> {
2309    let object = value.as_object()?;
2310    let expected = expected_query_id.to_string();
2311    if object.get("query_id")?.as_str()? != expected
2312        || !object.get("committed")?.as_bool()?
2313        || object.get("retryable")?.as_bool()?
2314        || !matches!(
2315            object.get("status")?.as_str()?,
2316            "committed"
2317                | "committed_with_error"
2318                | "partially_committed"
2319                | "cancelled_after_commit"
2320                | "deadline_after_commit"
2321        )
2322    {
2323        return None;
2324    }
2325    let committed_statements =
2326        usize::try_from(object.get("committed_statements")?.as_u64()?).ok()?;
2327    if committed_statements == 0 {
2328        return None;
2329    }
2330    let epoch_text = object.get("last_commit_epoch_text")?.as_str()?.to_owned();
2331    let numeric_epoch = object.get("last_commit_epoch")?.as_u64()?;
2332    let epoch = exact_epoch(Some(&epoch_text), Some(numeric_epoch)).ok()??;
2333    if numeric_epoch != epoch {
2334        return None;
2335    }
2336    let outcome = object.get("outcome")?.as_object()?;
2337    let outcome_epoch_text = outcome.get("last_commit_epoch_text")?.as_str()?;
2338    let outcome_numeric_epoch = outcome.get("last_commit_epoch")?.as_u64()?;
2339    let outcome_epoch =
2340        exact_epoch(Some(outcome_epoch_text), Some(outcome_numeric_epoch)).ok()??;
2341    if !outcome.get("committed")?.as_bool()?
2342        || usize::try_from(outcome.get("committed_statements")?.as_u64()?).ok()?
2343            != committed_statements
2344        || outcome_epoch_text != epoch_text
2345        || outcome_epoch != epoch
2346        || outcome_numeric_epoch != outcome_epoch
2347    {
2348        return None;
2349    }
2350    let error = object.get("error")?.as_object()?;
2351    if error.get("query_id")?.as_str()? != expected
2352        || !error.get("committed")?.as_bool()?
2353        || error.get("retryable")?.as_bool()?
2354    {
2355        return None;
2356    }
2357    Some(SqlReceiptCommitProof {
2358        epoch,
2359        epoch_text,
2360        committed_statements,
2361    })
2362}
2363
2364enum SqlReceiptDecodeError {
2365    KnownCommit(ClientError),
2366    Unknown(String),
2367}
2368
2369fn decode_remote_sql_receipt(
2370    bytes: &[u8],
2371    query_id: mongreldb_query::QueryId,
2372    expected_original_query_id: Option<mongreldb_query::QueryId>,
2373) -> Result<RemoteSqlReceipt, SqlReceiptDecodeError> {
2374    let value = strict_json_value(bytes)
2375        .map_err(|error| SqlReceiptDecodeError::Unknown(format!("invalid SQL receipt: {error}")))?;
2376    let proof = sql_receipt_commit_proof(&value, query_id, expected_original_query_id);
2377    let receipt = serde_json::from_value::<RemoteSqlReceipt>(value).map_err(|error| {
2378        proof.clone().map_or_else(
2379            || SqlReceiptDecodeError::Unknown(format!("invalid SQL receipt: {error}")),
2380            |proof| {
2381                SqlReceiptDecodeError::KnownCommit(committed_sql_receipt_decode_error(
2382                    query_id,
2383                    proof,
2384                    format!("SQL committed but its receipt was invalid: {error}"),
2385                ))
2386            },
2387        )
2388    })?;
2389    validate_remote_sql_receipt(receipt, query_id, expected_original_query_id).map_err(|error| {
2390        proof.map_or_else(
2391            || SqlReceiptDecodeError::Unknown(error.clone()),
2392            |proof| {
2393                SqlReceiptDecodeError::KnownCommit(committed_sql_receipt_decode_error(
2394                    query_id,
2395                    proof,
2396                    format!("SQL committed but its receipt was invalid: {error}"),
2397                ))
2398            },
2399        )
2400    })
2401}
2402
2403#[derive(Debug, Clone, Deserialize, PartialEq)]
2404#[serde(deny_unknown_fields)]
2405pub struct RemoteSqlPageLimits {
2406    pub rows: usize,
2407    pub bytes: usize,
2408    pub tokens: usize,
2409}
2410
2411#[derive(Debug, Clone, Deserialize, PartialEq)]
2412#[serde(deny_unknown_fields)]
2413pub struct RemoteSqlPageMetadata {
2414    pub offset: usize,
2415    pub row_count: usize,
2416    pub total_rows: usize,
2417    pub byte_count: usize,
2418    pub estimated_tokens: usize,
2419    pub limits: RemoteSqlPageLimits,
2420    pub projection: Vec<String>,
2421    pub expires_at_ms: u64,
2422    pub snapshot: String,
2423    pub token_estimate: String,
2424}
2425
2426#[derive(Debug, Clone, Deserialize, PartialEq)]
2427#[serde(deny_unknown_fields)]
2428pub struct RemoteSqlPage {
2429    pub status: String,
2430    pub rows: Vec<serde_json::Value>,
2431    pub next_cursor: Option<String>,
2432    pub page: RemoteSqlPageMetadata,
2433}
2434
2435pub struct RemoteSqlQueryHandle {
2436    query_id: mongreldb_query::QueryId,
2437    client: MongrelClient,
2438    result: std::thread::JoinHandle<ClientResult<Vec<RecordBatch>>>,
2439}
2440
2441impl RemoteSqlQueryHandle {
2442    pub fn id(&self) -> mongreldb_query::QueryId {
2443        self.query_id
2444    }
2445
2446    pub fn cancel(&self) -> ClientResult<RemoteCancelOutcome> {
2447        self.client.cancel_sql(self.query_id)
2448    }
2449
2450    pub fn status(&self) -> ClientResult<RemoteQueryStatus> {
2451        self.client.query_status(self.query_id)
2452    }
2453
2454    pub fn wait(self) -> ClientResult<Vec<RecordBatch>> {
2455        match self.result.join() {
2456            Ok(result) => result,
2457            Err(_) => Err(self
2458                .client
2459                .recover_after_transport_loss(self.query_id, "SQL worker panicked".into())),
2460        }
2461    }
2462}
2463
2464#[derive(Serialize)]
2465struct SqlReq {
2466    sql: String,
2467    #[serde(skip_serializing_if = "Option::is_none")]
2468    format: Option<&'static str>,
2469    query_id: mongreldb_query::QueryId,
2470    #[serde(skip_serializing_if = "Option::is_none")]
2471    timeout_ms: Option<u64>,
2472    #[serde(skip_serializing_if = "Option::is_none")]
2473    max_output_rows: Option<u64>,
2474    #[serde(skip_serializing_if = "Option::is_none")]
2475    max_output_bytes: Option<u64>,
2476    #[serde(skip_serializing_if = "Option::is_none")]
2477    idempotency_key: Option<String>,
2478    #[serde(skip_serializing_if = "Option::is_none")]
2479    pagination: Option<SqlPaginationReq>,
2480}
2481
2482#[derive(Debug, Clone, Serialize)]
2483struct SqlPaginationReq {
2484    page_size_rows: u64,
2485    projection: Vec<String>,
2486    #[serde(skip_serializing_if = "Option::is_none")]
2487    max_page_bytes: Option<u64>,
2488    #[serde(skip_serializing_if = "Option::is_none")]
2489    max_page_tokens: Option<u64>,
2490}
2491
2492#[derive(Deserialize)]
2493#[serde(deny_unknown_fields)]
2494struct CountResp {
2495    count: u64,
2496}
2497
2498#[derive(Deserialize)]
2499#[serde(deny_unknown_fields)]
2500struct TableIdResponse {
2501    table_id: u64,
2502    table_id_text: String,
2503}
2504
2505#[derive(Deserialize)]
2506#[serde(deny_unknown_fields)]
2507struct RowIdResponse {
2508    row_id: String,
2509}
2510
2511#[derive(Deserialize)]
2512#[serde(deny_unknown_fields)]
2513struct EpochResponse {
2514    epoch: u64,
2515    epoch_text: String,
2516}
2517
2518#[derive(Deserialize)]
2519#[serde(deny_unknown_fields)]
2520struct CommittedWriteResponse {
2521    status: String,
2522    epoch: u64,
2523    epoch_text: String,
2524}
2525
2526#[derive(Deserialize)]
2527#[serde(deny_unknown_fields)]
2528struct TriggerDropResponse {
2529    status: String,
2530    epoch: u64,
2531    epoch_text: String,
2532    dropped_trigger: mongreldb_core::StoredTrigger,
2533    resource_tables: Vec<TriggerTableBindingResponse>,
2534}
2535
2536#[derive(Deserialize)]
2537#[serde(deny_unknown_fields)]
2538struct TriggerTableBindingResponse {
2539    name: String,
2540    table_id: u64,
2541    schema_id: u64,
2542}
2543
2544#[derive(Debug, Clone, Deserialize)]
2545#[serde(deny_unknown_fields)]
2546pub struct HistoryRetention {
2547    pub history_retention_epochs: u64,
2548    pub earliest_retained_epoch: u64,
2549}
2550
2551/// Server-side schema metadata for one table (subset of the server's descriptor).
2552#[derive(Debug, Clone, Deserialize)]
2553#[serde(deny_unknown_fields)]
2554pub struct TableSchemaInfo {
2555    pub schema_id: u64,
2556    pub columns: Vec<ColumnMeta>,
2557    #[serde(default)]
2558    pub indexes: Vec<IndexMeta>,
2559    pub constraints: ConstraintMeta,
2560}
2561
2562#[derive(Debug, Clone, Deserialize)]
2563#[serde(deny_unknown_fields)]
2564pub struct IndexMeta {
2565    pub name: String,
2566    pub column_id: u16,
2567    pub kind: String,
2568    pub predicate: Option<String>,
2569    #[serde(default)]
2570    pub options: serde_json::Value,
2571}
2572
2573#[derive(Debug, Clone, Deserialize)]
2574#[serde(deny_unknown_fields)]
2575pub struct ColumnMeta {
2576    pub id: u16,
2577    pub name: String,
2578    pub ty: String,
2579    pub primary_key: bool,
2580    pub nullable: bool,
2581    pub auto_increment: bool,
2582}
2583
2584#[derive(Debug, Clone, Default, Deserialize)]
2585#[serde(deny_unknown_fields)]
2586pub struct ConstraintMeta {
2587    #[serde(default)]
2588    pub uniques: Vec<serde_json::Value>,
2589    #[serde(default)]
2590    pub foreign_keys: Vec<serde_json::Value>,
2591    #[serde(default)]
2592    pub checks: Vec<serde_json::Value>,
2593}
2594
2595// ── /kit/txn typed models ───────────────────────────────────────────────────
2596
2597/// A typed atomic batch request for `/kit/txn`.
2598#[derive(Debug, Clone, Serialize)]
2599pub struct KitTxnRequest {
2600    #[serde(skip_serializing_if = "Option::is_none")]
2601    pub idempotency_key: Option<String>,
2602    pub ops: Vec<KitOp>,
2603}
2604
2605impl KitTxnRequest {
2606    pub fn new(ops: Vec<KitOp>) -> Self {
2607        Self {
2608            idempotency_key: None,
2609            ops,
2610        }
2611    }
2612    pub fn with_idempotency_key(mut self, key: impl Into<String>) -> Self {
2613        self.idempotency_key = Some(key.into());
2614        self
2615    }
2616}
2617
2618/// One operation in a `/kit/txn` batch.
2619#[derive(Debug, Clone, Serialize)]
2620#[serde(rename_all = "snake_case")]
2621pub enum KitOp {
2622    Put {
2623        table: String,
2624        cells: Vec<serde_json::Value>,
2625        returning: bool,
2626    },
2627    Upsert {
2628        table: String,
2629        cells: Vec<serde_json::Value>,
2630        #[serde(skip_serializing_if = "Option::is_none")]
2631        update_cells: Option<Vec<serde_json::Value>>,
2632        returning: bool,
2633    },
2634    Delete {
2635        table: String,
2636        row_id: u64,
2637    },
2638    DeleteByPk {
2639        table: String,
2640        pk: serde_json::Value,
2641    },
2642}
2643
2644impl KitOp {
2645    pub fn put(table: impl Into<String>, cells: Vec<serde_json::Value>) -> Self {
2646        KitOp::Put {
2647            table: table.into(),
2648            cells,
2649            returning: false,
2650        }
2651    }
2652    pub fn put_returning(table: impl Into<String>, cells: Vec<serde_json::Value>) -> Self {
2653        KitOp::Put {
2654            table: table.into(),
2655            cells,
2656            returning: true,
2657        }
2658    }
2659    pub fn upsert(table: impl Into<String>, cells: Vec<serde_json::Value>) -> Self {
2660        KitOp::Upsert {
2661            table: table.into(),
2662            cells,
2663            update_cells: None,
2664            returning: false,
2665        }
2666    }
2667    pub fn delete_by_pk(table: impl Into<String>, pk: serde_json::Value) -> Self {
2668        KitOp::DeleteByPk {
2669            table: table.into(),
2670            pk,
2671        }
2672    }
2673}
2674
2675/// A typed per-op result from `/kit/txn`.
2676#[derive(Debug, Clone, Deserialize)]
2677#[serde(deny_unknown_fields, rename_all = "snake_case", tag = "kind")]
2678pub enum KitOpResult {
2679    Put {
2680        row_id: Option<String>,
2681        auto_inc: Option<i64>,
2682        #[serde(default)]
2683        row: Option<Vec<serde_json::Value>>,
2684    },
2685    Upsert {
2686        action: String,
2687        auto_inc: Option<i64>,
2688        #[serde(default)]
2689        row: Option<Vec<serde_json::Value>>,
2690    },
2691    Deleted,
2692    NotFound,
2693}
2694
2695#[derive(Debug, Clone, Deserialize)]
2696#[serde(deny_unknown_fields)]
2697pub struct KitTxnResponse {
2698    pub status: String,
2699    pub epoch: u64,
2700    #[serde(default)]
2701    pub epoch_text: Option<String>,
2702    pub results: Vec<KitOpResult>,
2703}
2704
2705fn validate_kit_txn_response(
2706    response: KitTxnResponse,
2707    request: &KitTxnRequest,
2708) -> ClientResult<KitTxnResponse> {
2709    if response.status != "committed" {
2710        return Err(ClientError::Decode(
2711            "Kit transaction success status is not committed".into(),
2712        ));
2713    }
2714    if response.epoch_text.is_none()
2715        || exact_kit_epoch(response.epoch_text.as_deref(), Some(response.epoch))
2716            .map_err(ClientError::Decode)?
2717            .is_none()
2718    {
2719        return Err(ClientError::Decode(
2720            "Kit transaction response has no exact epoch".into(),
2721        ));
2722    }
2723    validate_kit_results(&response.results, request)?;
2724    Ok(response)
2725}
2726
2727fn kit_txn_outcome_unknown(message: impl Into<String>) -> ClientError {
2728    ClientError::Kit {
2729        code: KitErrorCode::QueryOutcomeUnknown,
2730        message: message.into(),
2731        op_index: None,
2732        status: 0,
2733        committed: None,
2734        epoch: None,
2735        epoch_text: None,
2736        retryable: Some(false),
2737    }
2738}
2739
2740fn kit_txn_auth_error(status: u16) -> ClientError {
2741    let (code, message) = if status == 401 {
2742        (KitErrorCode::AuthRequired, "authentication required")
2743    } else {
2744        (KitErrorCode::PermissionDenied, "permission denied")
2745    };
2746    ClientError::Kit {
2747        code,
2748        message: message.into(),
2749        op_index: None,
2750        status,
2751        committed: Some(false),
2752        epoch: None,
2753        epoch_text: None,
2754        retryable: Some(false),
2755    }
2756}
2757
2758fn committed_kit_txn_decode_error(
2759    status: u16,
2760    epoch: u64,
2761    epoch_text: String,
2762    message: impl Into<String>,
2763) -> ClientError {
2764    ClientError::Kit {
2765        code: KitErrorCode::CommitOutcome,
2766        message: message.into(),
2767        op_index: None,
2768        status,
2769        committed: Some(true),
2770        epoch: Some(epoch),
2771        epoch_text: Some(epoch_text),
2772        retryable: Some(false),
2773    }
2774}
2775
2776fn decode_kit_txn_success(
2777    body: &[u8],
2778    request: &KitTxnRequest,
2779    status: u16,
2780) -> ClientResult<KitTxnResponse> {
2781    let value = strict_json_value(body).map_err(|error| {
2782        kit_txn_outcome_unknown(format!("invalid Kit transaction success response: {error}"))
2783    })?;
2784    let epoch = value.get("epoch").and_then(serde_json::Value::as_u64);
2785    let epoch_text = value
2786        .get("epoch_text")
2787        .and_then(serde_json::Value::as_str)
2788        .map(str::to_owned);
2789    let proven_commit = value.get("status").and_then(serde_json::Value::as_str)
2790        == Some("committed")
2791        && epoch.is_some()
2792        && epoch_text
2793            .as_deref()
2794            .is_some_and(|text| exact_kit_epoch(Some(text), epoch).ok().flatten() == epoch);
2795    let (Some(epoch), Some(epoch_text)) = (epoch, epoch_text) else {
2796        return Err(kit_txn_outcome_unknown(
2797            "Kit transaction success response does not prove a commit",
2798        ));
2799    };
2800    if !proven_commit {
2801        return Err(kit_txn_outcome_unknown(
2802            "Kit transaction success response does not prove a commit",
2803        ));
2804    }
2805    if let Err(error) = exact_json_object_fields(
2806        &value,
2807        &["status", "epoch", "epoch_text", "results"],
2808        &["status", "epoch", "epoch_text", "results"],
2809        "Kit transaction success response",
2810    ) {
2811        return Err(committed_kit_txn_decode_error(
2812            status,
2813            epoch,
2814            epoch_text,
2815            format!("transaction committed but its response was invalid: {error}"),
2816        ));
2817    }
2818    if !value["results"].is_array() {
2819        return Err(committed_kit_txn_decode_error(
2820            status,
2821            epoch,
2822            epoch_text,
2823            "transaction committed but its results were not an array",
2824        ));
2825    }
2826    let response = serde_json::from_value::<KitTxnResponse>(value).map_err(|error| {
2827        committed_kit_txn_decode_error(
2828            status,
2829            epoch,
2830            epoch_text.clone(),
2831            format!("transaction committed but its result could not be decoded: {error}"),
2832        )
2833    })?;
2834    validate_kit_txn_response(response, request).map_err(|error| {
2835        committed_kit_txn_decode_error(
2836            status,
2837            epoch,
2838            epoch_text,
2839            format!("transaction committed but its result was invalid: {error}"),
2840        )
2841    })
2842}
2843
2844fn validate_kit_results(results: &[KitOpResult], request: &KitTxnRequest) -> ClientResult<()> {
2845    if results.len() != request.ops.len() {
2846        return Err(ClientError::Decode(
2847            "Kit transaction result count does not match request".into(),
2848        ));
2849    }
2850    for (index, (operation, result)) in request.ops.iter().zip(results).enumerate() {
2851        let matches = match (operation, result) {
2852            (KitOp::Put { returning, .. }, KitOpResult::Put { row_id, row, .. }) => {
2853                row_id.is_none()
2854                    && row.is_some() == *returning
2855                    && row.as_deref().is_none_or(valid_flat_kit_cells)
2856            }
2857            (KitOp::Upsert { returning, .. }, KitOpResult::Upsert { action, row, .. }) => {
2858                matches!(action.as_str(), "inserted" | "updated" | "unchanged")
2859                    && row.is_some() == *returning
2860                    && row.as_deref().is_none_or(valid_flat_kit_cells)
2861            }
2862            (KitOp::Delete { .. }, KitOpResult::Deleted) => true,
2863            (KitOp::DeleteByPk { .. }, KitOpResult::Deleted | KitOpResult::NotFound) => true,
2864            _ => false,
2865        };
2866        if !matches {
2867            return Err(ClientError::Decode(format!(
2868                "Kit transaction result {index} does not match request operation"
2869            )));
2870        }
2871    }
2872    Ok(())
2873}
2874
2875fn valid_flat_kit_cells(cells: &[serde_json::Value]) -> bool {
2876    if !cells.len().is_multiple_of(2) {
2877        return false;
2878    }
2879    let mut column_ids = std::collections::HashSet::new();
2880    cells.chunks_exact(2).all(|pair| {
2881        pair[0]
2882            .as_u64()
2883            .and_then(|column_id| u16::try_from(column_id).ok())
2884            .is_some_and(|column_id| column_ids.insert(column_id))
2885    })
2886}
2887
2888// ── /kit/query typed models ─────────────────────────────────────────────────
2889
2890/// A native typed query (`POST /kit/query`). `conditions` are raw JSON objects
2891/// mirroring the daemon's condition variants, e.g. `{"pk": {"value": 2}}`,
2892/// `{"range": {"column_id": 2, "lo": 0, "hi": 100}}`,
2893/// `{"ann": {"column_id": 5, "query": [...], "k": 10}}`.
2894#[derive(Debug, Clone, Serialize)]
2895pub struct KitQueryRequest {
2896    pub table: String,
2897    #[serde(skip_serializing_if = "Vec::is_empty")]
2898    pub conditions: Vec<serde_json::Value>,
2899    #[serde(skip_serializing_if = "Option::is_none")]
2900    pub projection: Option<Vec<u16>>,
2901    #[serde(skip_serializing_if = "Option::is_none")]
2902    pub limit: Option<usize>,
2903    #[serde(skip_serializing_if = "Option::is_none")]
2904    pub offset: Option<usize>,
2905    #[serde(skip_serializing_if = "Option::is_none")]
2906    pub cursor: Option<String>,
2907}
2908
2909#[derive(Debug, Clone, Deserialize)]
2910#[serde(deny_unknown_fields)]
2911pub struct KitQueryResponse {
2912    pub rows: Vec<KitQueryRow>,
2913    pub truncated: bool,
2914    #[serde(default)]
2915    pub next_cursor: Option<String>,
2916}
2917
2918#[derive(Debug, Clone, Deserialize)]
2919#[serde(deny_unknown_fields)]
2920pub struct KitQueryRow {
2921    pub row_id: String,
2922    /// Flat `[col_id, val, col_id, val, …]` cells.
2923    pub cells: Vec<serde_json::Value>,
2924}
2925
2926fn validate_kit_query_response(
2927    response: KitQueryResponse,
2928    request: &KitQueryRequest,
2929) -> ClientResult<KitQueryResponse> {
2930    if response.truncated != response.next_cursor.is_some()
2931        || response
2932            .next_cursor
2933            .as_ref()
2934            .is_some_and(|cursor| cursor.is_empty() || cursor.len() > 2_048)
2935    {
2936        return Err(ClientError::Decode(
2937            "Kit query continuation cursor is inconsistent".into(),
2938        ));
2939    }
2940    let limit = request
2941        .limit
2942        .unwrap_or(mongreldb_core::query::MAX_FINAL_LIMIT);
2943    if response.rows.len() > limit {
2944        return Err(ClientError::Decode(
2945            "Kit query returned more rows than requested".into(),
2946        ));
2947    }
2948    let projection = request.projection.as_ref().map(|columns| {
2949        columns
2950            .iter()
2951            .copied()
2952            .collect::<std::collections::HashSet<_>>()
2953    });
2954    let mut row_ids = std::collections::HashSet::new();
2955    for row in &response.rows {
2956        if row
2957            .row_id
2958            .parse::<u64>()
2959            .ok()
2960            .is_none_or(|row_id| row_id.to_string() != row.row_id)
2961            || !row_ids.insert(row.row_id.as_str())
2962            || !valid_flat_kit_cells(&row.cells)
2963            || projection.as_ref().is_some_and(|projection| {
2964                row.cells.chunks_exact(2).any(|pair| {
2965                    pair[0]
2966                        .as_u64()
2967                        .and_then(|column_id| u16::try_from(column_id).ok())
2968                        .is_none_or(|column_id| !projection.contains(&column_id))
2969                })
2970            })
2971        {
2972            return Err(ClientError::Decode(
2973                "Kit query row identifier or cell layout is invalid".into(),
2974            ));
2975        }
2976    }
2977    Ok(response)
2978}
2979
2980/// Cooperative execution limits accepted by every Kit AI endpoint.
2981#[derive(Debug, Clone, Copy, Default, Serialize)]
2982pub struct AiExecutionOptions {
2983    #[serde(skip_serializing_if = "Option::is_none")]
2984    pub deadline_ms: Option<u64>,
2985    #[serde(skip_serializing_if = "Option::is_none")]
2986    pub max_work: Option<usize>,
2987}
2988
2989#[derive(Serialize)]
2990struct WithAiExecutionOptions<'a, T> {
2991    #[serde(flatten)]
2992    request: &'a T,
2993    #[serde(flatten)]
2994    options: &'a AiExecutionOptions,
2995}
2996
2997#[derive(Debug, Clone, Serialize)]
2998pub struct KitRetrieveRequest {
2999    pub table: String,
3000    pub retriever: serde_json::Value,
3001}
3002
3003#[derive(Debug, Clone, Deserialize)]
3004#[serde(deny_unknown_fields)]
3005pub struct KitRetrieveResponse {
3006    pub hits: Vec<KitRetrieverHit>,
3007}
3008
3009#[derive(Debug, Clone, Serialize)]
3010pub struct KitAnnRerankRequest {
3011    pub table: String,
3012    pub column_id: u16,
3013    pub query: Vec<f32>,
3014    pub candidate_k: usize,
3015    pub limit: usize,
3016    pub metric: KitVectorMetric,
3017}
3018
3019#[derive(Debug, Clone, Copy, Serialize)]
3020#[serde(rename_all = "snake_case")]
3021pub enum KitVectorMetric {
3022    Cosine,
3023    DotProduct,
3024    Euclidean,
3025}
3026
3027#[derive(Debug, Clone, Deserialize)]
3028#[serde(deny_unknown_fields)]
3029pub struct KitAnnRerankResponse {
3030    pub hits: Vec<KitAnnRerankHit>,
3031}
3032
3033#[derive(Debug, Clone, Deserialize)]
3034#[serde(deny_unknown_fields)]
3035pub struct KitAnnRerankHit {
3036    pub row_id: String,
3037    pub hamming_distance: u32,
3038    pub exact_score: f32,
3039}
3040
3041#[derive(Debug, Clone, Deserialize)]
3042#[serde(deny_unknown_fields)]
3043pub struct KitRetrieverHit {
3044    pub row_id: String,
3045    pub rank: usize,
3046    pub score: KitScore,
3047}
3048
3049#[derive(Debug, Clone, Deserialize)]
3050#[serde(deny_unknown_fields)]
3051pub struct KitScore {
3052    pub kind: String,
3053    pub value: f64,
3054}
3055
3056#[derive(Debug, Clone, Serialize)]
3057pub struct KitSetSimilarityRequest {
3058    pub table: String,
3059    pub column_id: u16,
3060    pub members: Vec<serde_json::Value>,
3061    pub candidate_k: usize,
3062    pub min_jaccard: f32,
3063    pub limit: usize,
3064}
3065
3066#[derive(Debug, Clone, Deserialize)]
3067#[serde(deny_unknown_fields)]
3068pub struct KitSetSimilarityResponse {
3069    pub hits: Vec<KitSetSimilarityHit>,
3070}
3071
3072#[derive(Debug, Clone, Deserialize)]
3073#[serde(deny_unknown_fields)]
3074pub struct KitSetSimilarityHit {
3075    pub row_id: String,
3076    pub estimated_jaccard: f32,
3077    pub exact_jaccard: f32,
3078}
3079
3080#[derive(Debug, Clone, Serialize)]
3081pub struct KitSearchRequest {
3082    pub table: String,
3083    #[serde(skip_serializing_if = "Vec::is_empty")]
3084    pub must: Vec<serde_json::Value>,
3085    pub retrievers: Vec<serde_json::Value>,
3086    pub fusion: serde_json::Value,
3087    #[serde(skip_serializing_if = "Option::is_none")]
3088    pub rerank: Option<serde_json::Value>,
3089    pub limit: usize,
3090    #[serde(skip_serializing_if = "Option::is_none")]
3091    pub projection: Option<Vec<u16>>,
3092    #[serde(skip_serializing_if = "Option::is_none")]
3093    pub deadline_ms: Option<u64>,
3094    #[serde(skip_serializing_if = "Option::is_none")]
3095    pub max_work: Option<usize>,
3096    #[serde(default)]
3097    pub explain: bool,
3098    #[serde(skip_serializing_if = "Option::is_none")]
3099    pub cursor: Option<String>,
3100}
3101
3102#[derive(Debug, Clone, Deserialize)]
3103#[serde(deny_unknown_fields)]
3104pub struct KitSearchResponse {
3105    pub hits: Vec<KitSearchHit>,
3106    #[serde(default)]
3107    pub trace: Option<serde_json::Value>,
3108    #[serde(default)]
3109    pub next_cursor: Option<String>,
3110}
3111
3112#[derive(Debug, Clone, Deserialize)]
3113#[serde(deny_unknown_fields)]
3114pub struct KitSearchHit {
3115    pub row_id: String,
3116    pub cells: Vec<serde_json::Value>,
3117    pub components: Vec<KitComponentScore>,
3118    pub fused_score: f64,
3119    #[serde(default)]
3120    pub exact_rerank_score: Option<f32>,
3121    #[serde(default)]
3122    pub final_score: Option<f64>,
3123    #[serde(default)]
3124    pub final_rank: Option<usize>,
3125}
3126
3127#[derive(Debug, Clone, Deserialize)]
3128#[serde(deny_unknown_fields)]
3129pub struct KitComponentScore {
3130    pub retriever_name: String,
3131    pub rank: usize,
3132    pub raw_score: KitScore,
3133    pub contribution: f64,
3134}
3135
3136#[derive(Debug, Clone, Serialize)]
3137pub struct ProcedureRequest {
3138    pub procedure: mongreldb_core::StoredProcedure,
3139}
3140
3141#[derive(Debug, Clone, Serialize)]
3142pub struct TriggerRequest {
3143    pub trigger: mongreldb_core::StoredTrigger,
3144    #[serde(skip_serializing_if = "Option::is_none")]
3145    pub idempotency_key: Option<String>,
3146}
3147
3148#[derive(Debug, Clone, Deserialize)]
3149#[serde(deny_unknown_fields)]
3150pub struct ProcedureResponse {
3151    #[serde(default)]
3152    pub status: String,
3153    pub procedure: mongreldb_core::StoredProcedure,
3154}
3155
3156#[derive(Debug, Clone, Deserialize)]
3157#[serde(deny_unknown_fields)]
3158pub struct TriggerResponse {
3159    #[serde(default)]
3160    pub status: Option<String>,
3161    pub trigger: mongreldb_core::StoredTrigger,
3162}
3163
3164#[derive(Debug, Clone, Deserialize)]
3165#[serde(deny_unknown_fields)]
3166pub struct ProceduresResponse {
3167    pub procedures: Vec<mongreldb_core::StoredProcedure>,
3168}
3169
3170#[derive(Debug, Clone, Deserialize)]
3171#[serde(deny_unknown_fields)]
3172pub struct TriggersResponse {
3173    pub triggers: Vec<mongreldb_core::StoredTrigger>,
3174}
3175
3176#[derive(Debug, Clone, Serialize)]
3177pub struct ProcedureCallRequest {
3178    #[serde(default)]
3179    pub args: serde_json::Map<String, serde_json::Value>,
3180    #[serde(skip_serializing_if = "Option::is_none")]
3181    pub idempotency_key: Option<String>,
3182}
3183
3184#[derive(Debug, Clone, Deserialize)]
3185#[serde(deny_unknown_fields)]
3186pub struct ProcedureCallResponse {
3187    pub status: String,
3188    pub committed: bool,
3189    #[serde(default)]
3190    pub epoch: Option<u64>,
3191    #[serde(default)]
3192    pub epoch_text: Option<String>,
3193    pub result: serde_json::Value,
3194}
3195
3196fn validate_procedure_call_response(
3197    response: ProcedureCallResponse,
3198) -> ClientResult<ProcedureCallResponse> {
3199    if response.status != "ok"
3200        || response.committed != response.epoch.is_some()
3201        || response.committed != response.epoch_text.is_some()
3202        || exact_kit_epoch(response.epoch_text.as_deref(), response.epoch)
3203            .map_err(ClientError::Decode)?
3204            != response.epoch
3205    {
3206        return Err(ClientError::Decode(
3207            "procedure call response has invalid durable metadata".into(),
3208        ));
3209    }
3210    Ok(response)
3211}
3212
3213fn procedure_definition_matches(
3214    actual: &mongreldb_core::StoredProcedure,
3215    requested: &mongreldb_core::StoredProcedure,
3216) -> bool {
3217    actual.mode == requested.mode
3218        && actual.params == requested.params
3219        && actual.body == requested.body
3220}
3221
3222fn trigger_definition_matches(
3223    actual: &mongreldb_core::StoredTrigger,
3224    requested: &mongreldb_core::StoredTrigger,
3225) -> bool {
3226    actual.target == requested.target
3227        && actual.timing == requested.timing
3228        && actual.event == requested.event
3229        && actual.update_of == requested.update_of
3230        && actual.target_columns == requested.target_columns
3231        && actual.when == requested.when
3232        && actual.program == requested.program
3233        && actual.enabled
3234}
3235
3236// Internal mirror of the server's error envelope.
3237#[derive(Debug, Deserialize)]
3238#[serde(deny_unknown_fields)]
3239struct KitErrorEnvelope {
3240    #[allow(dead_code)]
3241    status: String,
3242    #[serde(default)]
3243    committed: Option<bool>,
3244    #[serde(default)]
3245    epoch: Option<u64>,
3246    #[serde(default)]
3247    epoch_text: Option<String>,
3248    #[serde(default)]
3249    retryable: Option<bool>,
3250    #[serde(default)]
3251    results: Option<Vec<KitOpResult>>,
3252    error: KitErrorBody,
3253}
3254
3255#[derive(Debug, Deserialize)]
3256#[serde(deny_unknown_fields)]
3257struct KitErrorBody {
3258    code: String,
3259    message: String,
3260    #[serde(default)]
3261    op_index: Option<usize>,
3262}
3263
3264fn exact_json_object_fields(
3265    value: &serde_json::Value,
3266    allowed: &[&str],
3267    required: &[&str],
3268    context: &str,
3269) -> Result<(), String> {
3270    let object = value
3271        .as_object()
3272        .ok_or_else(|| format!("{context} is not an object"))?;
3273    if let Some(field) = object
3274        .keys()
3275        .find(|field| !allowed.contains(&field.as_str()))
3276    {
3277        return Err(format!("{context} contains unknown field {field:?}"));
3278    }
3279    if let Some(field) = required.iter().find(|field| !object.contains_key(**field)) {
3280        return Err(format!("{context} lacks required field {field:?}"));
3281    }
3282    Ok(())
3283}
3284
3285fn is_exact_query_not_found_response(
3286    body: &[u8],
3287    expected_query_id: mongreldb_query::QueryId,
3288) -> bool {
3289    const FIELDS: &[&str] = &[
3290        "query_id",
3291        "status",
3292        "terminal_state",
3293        "committed",
3294        "committed_statements",
3295        "last_commit_epoch",
3296        "last_commit_epoch_text",
3297        "first_commit_statement_index",
3298        "last_commit_statement_index",
3299        "completed_statements",
3300        "statement_index",
3301        "cancel_outcome",
3302        "cancellation_reason",
3303        "retryable",
3304        "server_state",
3305        "outcome",
3306        "error",
3307    ];
3308    const OUTCOME_FIELDS: &[&str] = &[
3309        "committed",
3310        "committed_statements",
3311        "last_commit_epoch",
3312        "last_commit_epoch_text",
3313        "first_commit_statement_index",
3314        "last_commit_statement_index",
3315        "completed_statements",
3316        "statement_index",
3317        "serialization",
3318    ];
3319    const NULL_FIELDS: &[&str] = &[
3320        "terminal_state",
3321        "committed",
3322        "committed_statements",
3323        "last_commit_epoch",
3324        "last_commit_epoch_text",
3325        "first_commit_statement_index",
3326        "last_commit_statement_index",
3327        "completed_statements",
3328        "statement_index",
3329        "cancellation_reason",
3330    ];
3331    let Ok(value) = strict_json_value(body) else {
3332        return false;
3333    };
3334    if exact_json_object_fields(&value, FIELDS, FIELDS, "query-not-found response").is_err() {
3335        return false;
3336    }
3337    let expected_query_id = expected_query_id.to_string();
3338    if value["query_id"].as_str() != Some(expected_query_id.as_str())
3339        || value["status"].as_str() != Some("unknown")
3340        || value["cancel_outcome"].as_str() != Some("not_found")
3341        || value["server_state"].as_str() != Some("not_found")
3342        || value["retryable"].as_bool() != Some(false)
3343        || NULL_FIELDS.iter().any(|field| !value[*field].is_null())
3344    {
3345        return false;
3346    }
3347    let outcome = &value["outcome"];
3348    if exact_json_object_fields(
3349        outcome,
3350        OUTCOME_FIELDS,
3351        OUTCOME_FIELDS,
3352        "query-not-found outcome",
3353    )
3354    .is_err()
3355        || outcome["serialization"].as_str() != Some("unknown")
3356        || OUTCOME_FIELDS
3357            .iter()
3358            .filter(|field| **field != "serialization")
3359            .any(|field| !outcome[*field].is_null())
3360    {
3361        return false;
3362    }
3363    let error = &value["error"];
3364    exact_json_object_fields(
3365        error,
3366        &["code", "message", "query_id", "committed", "retryable"],
3367        &["code", "message", "query_id", "committed", "retryable"],
3368        "query-not-found error",
3369    )
3370    .is_ok()
3371        && error["code"].as_str() == Some("QUERY_NOT_FOUND")
3372        && error["message"].as_str().is_some()
3373        && error["query_id"].as_str() == Some(expected_query_id.as_str())
3374        && error["committed"].is_null()
3375        && error["retryable"].as_bool() == Some(false)
3376}
3377
3378fn validate_kit_txn_error_json(value: &serde_json::Value) -> Result<(), String> {
3379    let error = value
3380        .get("error")
3381        .ok_or_else(|| "Kit transaction error response lacks error".to_owned())?;
3382    exact_json_object_fields(
3383        error,
3384        &["code", "message", "op_index"],
3385        &["code", "message"],
3386        "Kit transaction error",
3387    )?;
3388    let status = value
3389        .get("status")
3390        .and_then(serde_json::Value::as_str)
3391        .ok_or_else(|| "Kit transaction error response lacks status".to_owned())?;
3392    if matches!(status, "committed" | "outcome_unknown")
3393        && error
3394            .as_object()
3395            .is_some_and(|error| error.contains_key("op_index"))
3396    {
3397        return Err("durable Kit transaction error must omit op_index".into());
3398    }
3399    match status {
3400        "aborted" => {
3401            let enriched = value.get("committed").is_some() || value.get("retryable").is_some();
3402            if enriched {
3403                exact_json_object_fields(
3404                    value,
3405                    &["status", "committed", "retryable", "error"],
3406                    &["status", "committed", "retryable", "error"],
3407                    "Kit transaction error response",
3408                )?;
3409            } else {
3410                exact_json_object_fields(
3411                    value,
3412                    &["status", "error"],
3413                    &["status", "error"],
3414                    "Kit transaction error response",
3415                )?;
3416            }
3417        }
3418        "committed" => exact_json_object_fields(
3419            value,
3420            &[
3421                "status",
3422                "committed",
3423                "epoch",
3424                "epoch_text",
3425                "results",
3426                "retryable",
3427                "error",
3428            ],
3429            &[
3430                "status",
3431                "committed",
3432                "epoch",
3433                "epoch_text",
3434                "retryable",
3435                "error",
3436            ],
3437            "Kit transaction error response",
3438        )?,
3439        "outcome_unknown" => exact_json_object_fields(
3440            value,
3441            &[
3442                "status",
3443                "committed",
3444                "epoch",
3445                "epoch_text",
3446                "retryable",
3447                "error",
3448            ],
3449            &[
3450                "status",
3451                "committed",
3452                "epoch",
3453                "epoch_text",
3454                "retryable",
3455                "error",
3456            ],
3457            "Kit transaction error response",
3458        )?,
3459        _ => return Err("Kit transaction error response status is invalid".into()),
3460    }
3461    Ok(())
3462}
3463
3464fn proven_kit_commit_metadata(value: &serde_json::Value) -> Option<(u64, String)> {
3465    if value.get("status")?.as_str()? != "committed"
3466        || !value.get("committed")?.as_bool()?
3467        || value.get("retryable")?.as_bool()?
3468    {
3469        return None;
3470    }
3471    let error = value.get("error")?.as_object()?;
3472    if error.get("code")?.as_str()? != "COMMIT_OUTCOME"
3473        || error.get("message")?.as_str()?.is_empty()
3474        || error.get("op_index").is_some_and(|index| {
3475            !index.is_null()
3476                && index
3477                    .as_u64()
3478                    .and_then(|index| usize::try_from(index).ok())
3479                    .is_none()
3480        })
3481    {
3482        return None;
3483    }
3484    let epoch = value.get("epoch")?.as_u64()?;
3485    let epoch_text = value.get("epoch_text")?.as_str()?.to_owned();
3486    (exact_kit_epoch(Some(&epoch_text), Some(epoch))
3487        .ok()
3488        .flatten()
3489        == Some(epoch))
3490    .then_some((epoch, epoch_text))
3491}
3492
3493fn decode_http_error(status: u16, body: &[u8]) -> ClientError {
3494    let value = match strict_json_value(body) {
3495        Ok(value) => value,
3496        Err(error) => {
3497            if body
3498                .iter()
3499                .copied()
3500                .find(|byte| !byte.is_ascii_whitespace())
3501                .is_some_and(|byte| matches!(byte, b'{' | b'['))
3502            {
3503                return ClientError::Decode(format!("invalid HTTP error response: {error}"));
3504            }
3505            return ClientError::Http {
3506                status,
3507                body: format!("non-JSON error response ({} bytes)", body.len()),
3508            };
3509        }
3510    };
3511    if status == 404
3512        && value
3513            .get("query_id")
3514            .and_then(serde_json::Value::as_str)
3515            .and_then(|query_id| query_id.parse::<mongreldb_query::QueryId>().ok())
3516            .is_some_and(|query_id| is_exact_query_not_found_response(body, query_id))
3517    {
3518        return match serde_json::from_value::<RemoteQueryErrorResponse>(value) {
3519            Ok(response) => ClientError::Query {
3520                status,
3521                code: RemoteQueryErrorCode::QueryNotFound,
3522                message: response.error.message.clone(),
3523                response: Box::new(response),
3524            },
3525            Err(error) => ClientError::Decode(format!("invalid query-not-found response: {error}")),
3526        };
3527    }
3528    let queryless_sql_response = value.get("query_id").is_none()
3529        && value
3530            .get("error")
3531            .and_then(|error| error.get("query_id"))
3532            .is_none()
3533        && value.get("outcome").is_some();
3534    if queryless_sql_response {
3535        return match serde_json::from_value::<RemoteQueryErrorResponse>(value) {
3536            Ok(response) => match validate_queryless_sql_error(response) {
3537                Ok(response) => ClientError::Query {
3538                    status,
3539                    code: response.error.code.clone(),
3540                    message: response.error.message.clone(),
3541                    response: Box::new(response),
3542                },
3543                Err(error) => {
3544                    ClientError::Decode(format!("invalid queryless SQL error response: {error}"))
3545                }
3546            },
3547            Err(error) => {
3548                ClientError::Decode(format!("invalid queryless SQL error response: {error}"))
3549            }
3550        };
3551    }
3552    let query_response = value.get("query_id").is_some()
3553        || value.get("outcome").is_some()
3554        || value.get("server_state").is_some()
3555        || value
3556            .get("error")
3557            .and_then(|error| error.get("query_id"))
3558            .is_some();
3559    if query_response {
3560        return match serde_json::from_value::<RemoteQueryErrorResponse>(value) {
3561            Ok(response) => {
3562                let query_id = response
3563                    .query_id
3564                    .as_deref()
3565                    .or(response.error.query_id.as_deref())
3566                    .and_then(|query_id| query_id.parse::<mongreldb_query::QueryId>().ok());
3567                match query_id
3568                    .ok_or_else(|| "query error response lacks a valid query_id".to_owned())
3569                    .and_then(|query_id| validate_remote_query_error(response, query_id))
3570                {
3571                    Ok(response) => ClientError::Query {
3572                        status,
3573                        code: response.error.code.clone(),
3574                        message: response.error.message.clone(),
3575                        response: Box::new(response),
3576                    },
3577                    Err(error) => {
3578                        ClientError::Decode(format!("invalid query error response: {error}"))
3579                    }
3580                }
3581            }
3582            Err(error) => ClientError::Decode(format!("invalid query error response: {error}")),
3583        };
3584    }
3585    match serde_json::from_value::<KitErrorEnvelope>(value) {
3586        Ok(env) => match validate_kit_error_envelope(&env) {
3587            Ok(epoch) => ClientError::Kit {
3588                code: KitErrorCode::from_str(&env.error.code),
3589                message: env.error.message,
3590                op_index: env.error.op_index,
3591                status,
3592                committed: env.committed,
3593                epoch,
3594                epoch_text: env.epoch_text,
3595                retryable: env.retryable,
3596            },
3597            Err(message) => ClientError::Decode(message),
3598        },
3599        Err(error) => ClientError::Decode(format!("invalid Kit error response: {error}")),
3600    }
3601}
3602
3603fn validate_kit_error_envelope(env: &KitErrorEnvelope) -> Result<Option<u64>, String> {
3604    if env.status.is_empty()
3605        || env.error.code.is_empty()
3606        || env.error.message.is_empty()
3607        || env.retryable == Some(true)
3608            && !matches!(
3609                env.error.code.as_str(),
3610                "IDEMPOTENCY_STORE_FULL" | "IDEMPOTENCY_STORE_UNAVAILABLE"
3611            )
3612    {
3613        return Err("invalid Kit error metadata".into());
3614    }
3615    let epoch = exact_kit_epoch(env.epoch_text.as_deref(), env.epoch)?;
3616    match env.error.code.as_str() {
3617        "COMMIT_OUTCOME" => {
3618            if env.status != "committed"
3619                || env.committed != Some(true)
3620                || epoch.is_none()
3621                || env.epoch_text.is_none()
3622                || env.retryable != Some(false)
3623                || env.error.op_index.is_some()
3624            {
3625                return Err("invalid committed Kit error metadata".into());
3626            }
3627        }
3628        "QUERY_OUTCOME_UNKNOWN" => {
3629            if env.status != "outcome_unknown"
3630                || env.committed.is_some()
3631                || env.results.is_some()
3632                || env.retryable != Some(false)
3633                || env.error.op_index.is_some()
3634            {
3635                return Err("invalid unknown Kit outcome metadata".into());
3636            }
3637        }
3638        _ => {
3639            let metadata_valid = match env.status.as_str() {
3640                "aborted" => matches!(
3641                    (env.committed, env.retryable),
3642                    (None, None) | (Some(false), Some(_))
3643                ),
3644                "error" => env.committed.is_none() && env.retryable.is_none(),
3645                _ => false,
3646            };
3647            if !metadata_valid || epoch.is_some() || env.results.is_some() {
3648                return Err("invalid aborted Kit error metadata".into());
3649            }
3650        }
3651    }
3652    Ok(epoch)
3653}
3654
3655fn decode_kit_txn_http_error(status: u16, body: &[u8], request: &KitTxnRequest) -> ClientError {
3656    let value = match strict_json_value(body) {
3657        Ok(value) => value,
3658        Err(error) => {
3659            return kit_txn_outcome_unknown(format!(
3660                "invalid Kit transaction error response: {error}"
3661            ))
3662        }
3663    };
3664    let proven_commit = proven_kit_commit_metadata(&value);
3665    if let Err(error) = validate_kit_txn_error_json(&value) {
3666        if let Some((epoch, epoch_text)) = proven_commit.as_ref() {
3667            return committed_kit_txn_decode_error(status, *epoch, epoch_text.clone(), error);
3668        }
3669        return kit_txn_outcome_unknown(error);
3670    }
3671    let env = match serde_json::from_value::<KitErrorEnvelope>(value) {
3672        Ok(env) => env,
3673        Err(error) => {
3674            if let Some((epoch, epoch_text)) = proven_commit {
3675                return committed_kit_txn_decode_error(
3676                    status,
3677                    epoch,
3678                    epoch_text,
3679                    format!(
3680                        "transaction committed but its error result could not be decoded: {error}"
3681                    ),
3682                );
3683            }
3684            return kit_txn_outcome_unknown(format!(
3685                "invalid Kit transaction error response: {error}"
3686            ));
3687        }
3688    };
3689    let epoch = match validate_kit_error_envelope(&env) {
3690        Ok(epoch) => epoch,
3691        Err(error) => {
3692            if let Some((epoch, epoch_text)) = proven_commit {
3693                return committed_kit_txn_decode_error(status, epoch, epoch_text, error);
3694            }
3695            return kit_txn_outcome_unknown(error);
3696        }
3697    };
3698    if env
3699        .error
3700        .op_index
3701        .is_some_and(|op_index| op_index >= request.ops.len())
3702    {
3703        return kit_txn_outcome_unknown("Kit transaction error op_index is out of bounds");
3704    }
3705    if env.committed == Some(true) {
3706        if let Some(results) = env.results.as_deref() {
3707            if let Err(error) = validate_kit_results(results, request) {
3708                let (Some(epoch), Some(epoch_text)) = (epoch, env.epoch_text.clone()) else {
3709                    return kit_txn_outcome_unknown(
3710                        "committed Kit transaction response lost validated epoch metadata",
3711                    );
3712                };
3713                return committed_kit_txn_decode_error(
3714                    status,
3715                    epoch,
3716                    epoch_text,
3717                    format!("transaction committed but its result was invalid: {error}"),
3718                );
3719            }
3720        }
3721    }
3722    ClientError::Kit {
3723        code: KitErrorCode::from_str(&env.error.code),
3724        message: env.error.message,
3725        op_index: env.error.op_index,
3726        status,
3727        committed: env.committed,
3728        epoch,
3729        epoch_text: env.epoch_text,
3730        retryable: env.retryable,
3731    }
3732}
3733
3734fn decode_sql_http_error(
3735    status: u16,
3736    body: &[u8],
3737    expected_query_id: mongreldb_query::QueryId,
3738) -> Result<ClientError, String> {
3739    let value =
3740        strict_json_value(body).map_err(|error| format!("invalid SQL error response: {error}"))?;
3741    if status == 404 && is_exact_query_not_found_response(body, expected_query_id) {
3742        let response = serde_json::from_value::<RemoteQueryErrorResponse>(value)
3743            .map_err(|error| format!("invalid query-not-found response: {error}"))?;
3744        return Ok(ClientError::Query {
3745            status,
3746            code: RemoteQueryErrorCode::QueryNotFound,
3747            message: response.error.message.clone(),
3748            response: Box::new(response),
3749        });
3750    }
3751    let proof = sql_error_commit_proof(&value, expected_query_id);
3752    let response = match serde_json::from_value::<RemoteQueryErrorResponse>(value) {
3753        Ok(response) => response,
3754        Err(error) => {
3755            if let Some(proof) = proof.clone() {
3756                return Ok(committed_sql_receipt_decode_error(
3757                    expected_query_id,
3758                    proof,
3759                    format!("SQL committed but its error response was invalid: {error}"),
3760                ));
3761            }
3762            return Err(format!("invalid SQL error response: {error}"));
3763        }
3764    };
3765    let response = match validate_remote_query_error(response, expected_query_id) {
3766        Ok(response) => response,
3767        Err(error) => {
3768            if let Some(proof) = proof {
3769                return Ok(committed_sql_receipt_decode_error(
3770                    expected_query_id,
3771                    proof,
3772                    format!("SQL committed but its error response was invalid: {error}"),
3773                ));
3774            }
3775            return Err(error);
3776        }
3777    };
3778    Ok(ClientError::Query {
3779        status,
3780        code: response.error.code.clone(),
3781        message: response.error.message.clone(),
3782        response: Box::new(response),
3783    })
3784}
3785
3786fn validate_sql_query_id_header(
3787    headers: &reqwest::header::HeaderMap,
3788    expected_query_id: mongreldb_query::QueryId,
3789) -> Result<(), String> {
3790    let value = headers
3791        .get("x-mongreldb-query-id")
3792        .ok_or_else(|| "SQL response is missing x-mongreldb-query-id".to_owned())?
3793        .to_str()
3794        .map_err(|_| "SQL response has a non-UTF-8 x-mongreldb-query-id".to_owned())?;
3795    if value != expected_query_id.to_string() {
3796        return Err("SQL response x-mongreldb-query-id does not match the request".into());
3797    }
3798    Ok(())
3799}
3800
3801fn client_error_proves_commit(error: &ClientError) -> bool {
3802    matches!(
3803        error,
3804        ClientError::Query { response, .. }
3805            if response.committed == Some(true)
3806                && response.committed_statements.is_some_and(|count| count > 0)
3807                && response.last_commit_epoch.is_some()
3808                && response.last_commit_epoch_text.is_some()
3809    )
3810}
3811
3812fn exact_kit_epoch(text: Option<&str>, numeric: Option<u64>) -> Result<Option<u64>, String> {
3813    let Some(text) = text else {
3814        return Ok(numeric);
3815    };
3816    let exact = text
3817        .parse::<u64>()
3818        .map_err(|_| "epoch_text is not an unsigned integer".to_owned())?;
3819    if exact.to_string() != text {
3820        return Err("epoch_text is not canonical".into());
3821    }
3822    if numeric.is_some_and(|numeric| numeric != exact) {
3823        return Err("epoch and epoch_text disagree".into());
3824    }
3825    Ok(Some(exact))
3826}
3827
3828fn exact_required_u64(field: &str, numeric: u64, text: &str) -> ClientResult<u64> {
3829    let exact = text
3830        .parse::<u64>()
3831        .map_err(|_| ClientError::Decode(format!("{field}_text is not an unsigned integer")))?;
3832    if exact.to_string() != text {
3833        return Err(ClientError::Decode(format!(
3834            "{field}_text is not canonical"
3835        )));
3836    }
3837    if exact != numeric {
3838        return Err(ClientError::Decode(format!(
3839            "{field} and {field}_text disagree"
3840        )));
3841    }
3842    Ok(exact)
3843}
3844
3845fn parse_required_u64_text(field: &str, text: &str) -> ClientResult<u64> {
3846    let value = text
3847        .parse::<u64>()
3848        .map_err(|_| ClientError::Decode(format!("{field} is not an unsigned integer")))?;
3849    if value.to_string() != text {
3850        return Err(ClientError::Decode(format!("{field} is not canonical")));
3851    }
3852    Ok(value)
3853}
3854
3855fn decode_required_json<T: serde::de::DeserializeOwned>(
3856    response: reqwest::blocking::Response,
3857    context: &str,
3858) -> ClientResult<T> {
3859    decode_blocking_json(response, MAX_CONTROL_RESPONSE_BYTES, context)
3860}
3861
3862fn write_outcome_unknown(context: &str, error: impl std::fmt::Display) -> ClientError {
3863    kit_txn_outcome_unknown(format!(
3864        "{context} outcome is unknown because its response could not be confirmed: {error}"
3865    ))
3866}
3867
3868fn decode_write_json<T: serde::de::DeserializeOwned>(
3869    response: reqwest::blocking::Response,
3870    context: &str,
3871) -> ClientResult<T> {
3872    decode_required_json(response, context).map_err(|error| write_outcome_unknown(context, error))
3873}
3874
3875fn validate_committed_write(response: &CommittedWriteResponse, context: &str) -> ClientResult<u64> {
3876    if response.status != "committed" {
3877        return Err(write_outcome_unknown(
3878            context,
3879            "success response does not prove a commit",
3880        ));
3881    }
3882    exact_required_u64("epoch", response.epoch, &response.epoch_text)
3883        .map_err(|error| write_outcome_unknown(context, error))
3884}
3885
3886fn validate_trigger_drop_response(
3887    response: &TriggerDropResponse,
3888    expected_name: &str,
3889) -> ClientResult<()> {
3890    validate_committed_write(
3891        &CommittedWriteResponse {
3892            status: response.status.clone(),
3893            epoch: response.epoch,
3894            epoch_text: response.epoch_text.clone(),
3895        },
3896        "drop trigger",
3897    )?;
3898    let mut names = std::collections::HashSet::new();
3899    let mut identities = std::collections::HashSet::new();
3900    if response.dropped_trigger.name != expected_name
3901        || response.resource_tables.iter().any(|binding| {
3902            binding.name.is_empty()
3903                || !names.insert(&binding.name)
3904                || !identities.insert((binding.table_id, binding.schema_id))
3905        })
3906    {
3907        return Err(write_outcome_unknown(
3908            "drop trigger",
3909            "success response has an invalid resource binding",
3910        ));
3911    }
3912    Ok(())
3913}
3914
3915fn capability_unsupported(message: impl Into<String>) -> ClientError {
3916    let message = message.into();
3917    let response = RemoteQueryErrorResponse {
3918        query_id: None,
3919        status: "failed_before_commit".into(),
3920        terminal_state: Some("failed_before_commit".into()),
3921        committed: Some(false),
3922        committed_statements: Some(0),
3923        last_commit_epoch: None,
3924        last_commit_epoch_text: None,
3925        first_commit_statement_index: None,
3926        last_commit_statement_index: None,
3927        completed_statements: Some(0),
3928        statement_index: Some(0),
3929        cancel_outcome: None,
3930        cancellation_reason: None,
3931        retryable: false,
3932        server_state: None,
3933        outcome: RemoteQueryOutcome::default(),
3934        error: RemoteQueryErrorBody {
3935            code: RemoteQueryErrorCode::CapabilityUnsupported,
3936            message: message.clone(),
3937            query_id: None,
3938            committed: Some(false),
3939            retryable: false,
3940        },
3941    };
3942    ClientError::Query {
3943        status: 0,
3944        code: RemoteQueryErrorCode::CapabilityUnsupported,
3945        message,
3946        response: Box::new(response),
3947    }
3948}
3949
3950fn client_serialization_error(
3951    query_id: Option<mongreldb_query::QueryId>,
3952    message: impl Into<String>,
3953) -> ClientError {
3954    let message = message.into();
3955    let query_id = query_id.map(|query_id| query_id.to_string());
3956    let response = RemoteQueryErrorResponse {
3957        query_id: query_id.clone(),
3958        status: "failed_before_commit".into(),
3959        terminal_state: Some("failed_before_commit".into()),
3960        committed: Some(false),
3961        committed_statements: Some(0),
3962        last_commit_epoch: None,
3963        last_commit_epoch_text: None,
3964        first_commit_statement_index: None,
3965        last_commit_statement_index: None,
3966        completed_statements: Some(0),
3967        statement_index: Some(0),
3968        cancel_outcome: None,
3969        cancellation_reason: None,
3970        retryable: false,
3971        server_state: Some("failed".into()),
3972        outcome: RemoteQueryOutcome {
3973            committed: Some(false),
3974            committed_statements: Some(0),
3975            completed_statements: Some(0),
3976            statement_index: Some(0),
3977            serialization: "failed".into(),
3978            ..RemoteQueryOutcome::default()
3979        },
3980        error: RemoteQueryErrorBody {
3981            code: RemoteQueryErrorCode::SerializationFailed,
3982            message: message.clone(),
3983            query_id,
3984            committed: Some(false),
3985            retryable: false,
3986        },
3987    };
3988    ClientError::Query {
3989        status: 0,
3990        code: RemoteQueryErrorCode::SerializationFailed,
3991        message,
3992        response: Box::new(response),
3993    }
3994}
3995
3996struct IdempotentAttemptError {
3997    error: ClientError,
3998    replay: bool,
3999}
4000
4001impl IdempotentAttemptError {
4002    fn final_error(error: ClientError) -> Self {
4003        Self {
4004            error,
4005            replay: false,
4006        }
4007    }
4008}
4009
4010fn fresh_query_id(previous: mongreldb_query::QueryId) -> ClientResult<mongreldb_query::QueryId> {
4011    let query_id = mongreldb_query::QueryId::random()
4012        .map_err(|error| ClientError::Transport(error.to_string()))?;
4013    if query_id == previous {
4014        return Err(ClientError::Transport(
4015            "generated duplicate SQL query ID".into(),
4016        ));
4017    }
4018    Ok(query_id)
4019}
4020
4021fn max_known(left: Option<usize>, right: Option<usize>) -> Option<usize> {
4022    left.into_iter().chain(right).max()
4023}
4024
4025fn recovered_query_error(status: RemoteQueryStatus, message: String) -> ClientError {
4026    let committed = status.durably_committed();
4027    let code = status
4028        .terminal_error
4029        .as_ref()
4030        .map(|error| error.code.clone())
4031        .unwrap_or_else(|| {
4032            if committed == Some(true) {
4033                RemoteQueryErrorCode::CommitOutcome
4034            } else {
4035                RemoteQueryErrorCode::SerializationFailed
4036            }
4037        });
4038    let response = RemoteQueryErrorResponse {
4039        query_id: Some(status.query_id.clone()),
4040        status: status.status.clone(),
4041        terminal_state: status.terminal_state.clone(),
4042        committed,
4043        committed_statements: max_known(
4044            status.committed_statements,
4045            status.outcome.committed_statements,
4046        ),
4047        last_commit_epoch: status
4048            .last_commit_epoch
4049            .or(status.outcome.last_commit_epoch),
4050        last_commit_epoch_text: status
4051            .last_commit_epoch_text
4052            .clone()
4053            .or_else(|| status.outcome.last_commit_epoch_text.clone()),
4054        first_commit_statement_index: status
4055            .first_commit_statement_index
4056            .or(status.outcome.first_commit_statement_index),
4057        last_commit_statement_index: status
4058            .last_commit_statement_index
4059            .or(status.outcome.last_commit_statement_index),
4060        completed_statements: max_known(
4061            status.completed_statements,
4062            status.outcome.completed_statements,
4063        ),
4064        statement_index: max_known(status.statement_index, status.outcome.statement_index),
4065        cancel_outcome: status.cancel_outcome,
4066        cancellation_reason: (!status.cancellation_reason.is_empty())
4067            .then_some(status.cancellation_reason.clone()),
4068        retryable: status.retryable,
4069        server_state: Some(status.server_state_or_state().to_owned()),
4070        outcome: status.outcome,
4071        error: RemoteQueryErrorBody {
4072            code: code.clone(),
4073            message: message.clone(),
4074            query_id: Some(status.query_id),
4075            committed,
4076            retryable: status.retryable,
4077        },
4078    };
4079    ClientError::Query {
4080        status: 0,
4081        code,
4082        message,
4083        response: Box::new(response),
4084    }
4085}
4086
4087fn recovery_status_is_decisive(status: &RemoteQueryStatus) -> bool {
4088    status.durably_committed() == Some(true)
4089        || status.is_terminal()
4090            && (status.durably_committed().is_some() || status.terminal_error.is_some())
4091}
4092
4093impl MongrelClient {
4094    pub fn builder(url: impl AsRef<str>) -> MongrelClientBuilder {
4095        let base_url = sanitized_base_url(url.as_ref());
4096        MongrelClientBuilder {
4097            invalid_base_url: base_url.is_none(),
4098            base_url: base_url.unwrap_or_default(),
4099            authorization: None,
4100            invalid_authorization: false,
4101            connect_timeout: None,
4102            request_timeout: None,
4103            pool_idle_timeout: None,
4104        }
4105    }
4106
4107    pub fn new(url: &str) -> ClientResult<Self> {
4108        Self::builder(url).build()
4109    }
4110
4111    pub fn with_options(url: impl AsRef<str>, options: RemoteOptions) -> ClientResult<Self> {
4112        let mut builder = Self::builder(url);
4113        if let Some(timeout) = options.transport_timeout {
4114            builder = builder.request_timeout(timeout);
4115        }
4116        builder = match options.auth {
4117            Some(RemoteAuth::Bearer(token)) => builder.bearer_token(token.expose_secret()),
4118            Some(RemoteAuth::Basic { username, password }) => {
4119                builder.basic_auth(username, password.expose_secret())
4120            }
4121            None => builder,
4122        };
4123        builder.build()
4124    }
4125
4126    fn url(&self, path: &str) -> String {
4127        format!("{}{path}", self.base_url)
4128    }
4129
4130    fn url_segments(&self, segments: &[&str]) -> ClientResult<String> {
4131        url_with_segments(&self.base_url, segments)
4132    }
4133
4134    /// Send a request and, on non-2xx, decode a typed [`ClientError::Kit`] when
4135    /// the body is a `KitErrorEnvelope`, else a plain [`ClientError::Http`].
4136    fn check(
4137        &self,
4138        resp: reqwest::blocking::Response,
4139    ) -> ClientResult<reqwest::blocking::Response> {
4140        let status = resp.status();
4141        if status.is_success() {
4142            return Ok(resp);
4143        }
4144        let status_u16 = status.as_u16();
4145        let body = bounded_blocking_bytes(resp, MAX_CONTROL_RESPONSE_BYTES).map_err(|error| {
4146            ClientError::Decode(format!("invalid HTTP error response: {error}"))
4147        })?;
4148        Err(decode_http_error(status_u16, &body))
4149    }
4150
4151    fn write_response(
4152        &self,
4153        response: Result<reqwest::blocking::Response, reqwest::Error>,
4154        context: &str,
4155    ) -> ClientResult<reqwest::blocking::Response> {
4156        let response = response.map_err(|error| write_outcome_unknown(context, error))?;
4157        let status = response.status();
4158        if status.is_success() {
4159            return Ok(response);
4160        }
4161        let body = bounded_blocking_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
4162            .map_err(|error| write_outcome_unknown(context, error))?;
4163        let error = decode_http_error(status.as_u16(), &body);
4164        match error {
4165            ClientError::Decode(message) => Err(write_outcome_unknown(context, message)),
4166            error => Err(error),
4167        }
4168    }
4169
4170    fn check_sql_response(
4171        &self,
4172        response: reqwest::blocking::Response,
4173        query_id: mongreldb_query::QueryId,
4174    ) -> ClientResult<reqwest::blocking::Response> {
4175        let status = response.status();
4176        let pre_handler_auth = matches!(
4177            status,
4178            reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN
4179        );
4180        if status.is_success() {
4181            return Ok(response);
4182        }
4183        let header_error = (!pre_handler_auth)
4184            .then(|| validate_sql_query_id_header(response.headers(), query_id).err())
4185            .flatten();
4186        let status = status.as_u16();
4187        let body =
4188            bounded_blocking_bytes(response, MAX_CONTROL_RESPONSE_BYTES).map_err(|error| {
4189                if pre_handler_auth {
4190                    ClientError::Http {
4191                        status,
4192                        body: format!("unreadable authentication error response: {error}"),
4193                    }
4194                } else {
4195                    self.recover_after_transport_loss(query_id, error.to_string())
4196                }
4197            })?;
4198        if pre_handler_auth {
4199            return Err(decode_http_error(status, &body));
4200        }
4201        match decode_sql_http_error(status, &body, query_id) {
4202            Ok(error) if client_error_proves_commit(&error) => Err(error),
4203            Ok(error) if header_error.is_none() => Err(error),
4204            Ok(_) => Err(self.recover_after_transport_loss(
4205                query_id,
4206                header_error.unwrap_or_else(|| "SQL response identity is invalid".into()),
4207            )),
4208            Err(error) => Err(self.recover_after_transport_loss(query_id, error)),
4209        }
4210    }
4211
4212    fn check_query_status_response(
4213        &self,
4214        response: reqwest::blocking::Response,
4215        query_id: mongreldb_query::QueryId,
4216    ) -> ClientResult<reqwest::blocking::Response> {
4217        let status = response.status();
4218        if status.is_success() {
4219            return Ok(response);
4220        }
4221        let status = status.as_u16();
4222        let body = bounded_blocking_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
4223            .map_err(ClientError::Decode)?;
4224        if matches!(status, 401 | 403) {
4225            return Err(decode_http_error(status, &body));
4226        }
4227        match decode_sql_http_error(status, &body, query_id) {
4228            Ok(error) => Err(error),
4229            Err(error) => Err(ClientError::Decode(error)),
4230        }
4231    }
4232
4233    pub fn health(&self) -> ClientResult<String> {
4234        let resp = self.client.get(self.url("/health")).send()?;
4235        let bytes = bounded_blocking_bytes(self.check(resp)?, MAX_CONTROL_RESPONSE_BYTES)
4236            .map_err(ClientError::Transport)?;
4237        String::from_utf8(bytes)
4238            .map_err(|_| ClientError::Decode("invalid health response: non-UTF-8 body".into()))
4239    }
4240
4241    pub fn capabilities(&self) -> ClientResult<ServerCapabilities> {
4242        let response = self.client.get(self.url("/capabilities")).send()?;
4243        decode_blocking_json(
4244            self.check(response)?,
4245            MAX_CONTROL_RESPONSE_BYTES,
4246            "capabilities",
4247        )
4248    }
4249
4250    pub fn sql_cancellation_capabilities(&self) -> ClientResult<SqlCancellationCapabilities> {
4251        let capabilities = self.capabilities()?.sql_cancellation;
4252        if capabilities.version < 2
4253            || !capabilities.client_query_ids
4254            || !capabilities.cancel_endpoint
4255            || !capabilities.query_status
4256            || !capabilities.pre_registration_cancel
4257        {
4258            return Err(capability_unsupported(
4259                "server does not support SQL cancellation capability version 2",
4260            ));
4261        }
4262        Ok(capabilities)
4263    }
4264
4265    pub fn sql_idempotency_capabilities(&self) -> ClientResult<SqlIdempotencyCapabilities> {
4266        let capabilities = self
4267            .capabilities()?
4268            .sql_idempotency
4269            .ok_or_else(|| capability_unsupported("server does not advertise SQL idempotency"))?;
4270        if capabilities.version < 1
4271            || !capabilities.durable_pre_execution_intent
4272            || !capabilities.replay_committed_receipt
4273            || !capabilities.indeterminate_never_reexecutes
4274        {
4275            return Err(capability_unsupported(
4276                "server does not support safe SQL idempotency capability version 1",
4277            ));
4278        }
4279        Ok(capabilities)
4280    }
4281
4282    pub fn sql_pagination_capabilities(&self) -> ClientResult<SqlPaginationCapabilities> {
4283        let capabilities = self
4284            .capabilities()?
4285            .sql_pagination
4286            .ok_or_else(|| capability_unsupported("server does not advertise SQL pagination"))?;
4287        if capabilities.version < 1
4288            || capabilities.continuation_endpoint != "/sql/continue"
4289            || !capabilities.retained_snapshot
4290            || !capabilities.projection_required
4291            || !capabilities.byte_and_token_hints
4292        {
4293            return Err(capability_unsupported(
4294                "server does not support SQL pagination capability version 1",
4295            ));
4296        }
4297        Ok(capabilities)
4298    }
4299
4300    pub fn set_history_retention_epochs(&self, epochs: u64) -> ClientResult<HistoryRetention> {
4301        let resp = self
4302            .client
4303            .put(self.url("/history/retention"))
4304            .json(&serde_json::json!({"history_retention_epochs": epochs}))
4305            .send()?;
4306        decode_blocking_json(
4307            self.check(resp)?,
4308            MAX_CONTROL_RESPONSE_BYTES,
4309            "history retention",
4310        )
4311    }
4312
4313    pub fn history_retention_epochs(&self) -> ClientResult<u64> {
4314        Ok(self.history_retention()?.history_retention_epochs)
4315    }
4316
4317    pub fn earliest_retained_epoch(&self) -> ClientResult<u64> {
4318        Ok(self.history_retention()?.earliest_retained_epoch)
4319    }
4320
4321    fn history_retention(&self) -> ClientResult<HistoryRetention> {
4322        let resp = self.client.get(self.url("/history/retention")).send()?;
4323        decode_blocking_json(
4324            self.check(resp)?,
4325            MAX_CONTROL_RESPONSE_BYTES,
4326            "history retention",
4327        )
4328    }
4329
4330    // ── Table management ──
4331
4332    pub fn list_tables(&self) -> ClientResult<Vec<String>> {
4333        let resp = self.client.get(self.url("/tables")).send()?;
4334        decode_blocking_json(self.check(resp)?, MAX_CONTROL_RESPONSE_BYTES, "table list")
4335    }
4336
4337    pub fn create_table(&self, name: &str, columns: Vec<ColumnDefJson>) -> ClientResult<u64> {
4338        let response = self.write_response(
4339            self.client
4340                .post(self.url("/tables"))
4341                .json(&serde_json::json!({ "name": name, "columns": columns }))
4342                .send(),
4343            "create table",
4344        )?;
4345        let response: TableIdResponse = decode_write_json(response, "create table")?;
4346        exact_required_u64("table_id", response.table_id, &response.table_id_text)
4347            .map_err(|error| write_outcome_unknown("create table", error))
4348    }
4349
4350    pub fn drop_table(&self, name: &str) -> ClientResult<()> {
4351        let response = self.write_response(
4352            self.client
4353                .delete(self.url_segments(&["tables", name])?)
4354                .send(),
4355            "drop table",
4356        )?;
4357        let response: CommittedWriteResponse = decode_write_json(response, "drop table")?;
4358        validate_committed_write(&response, "drop table")?;
4359        Ok(())
4360    }
4361
4362    // ── Table-qualified operations ──
4363
4364    pub fn count(&self, table: &str) -> ClientResult<u64> {
4365        let resp = self
4366            .client
4367            .get(self.url_segments(&["tables", table, "count"])?)
4368            .send()?;
4369        let resp = self.check(resp)?;
4370        let cr: CountResp = decode_blocking_json(resp, MAX_CONTROL_RESPONSE_BYTES, "count")?;
4371        Ok(cr.count)
4372    }
4373
4374    pub fn put(&self, table: &str, row: Vec<(u16, mongreldb_core::Value)>) -> ClientResult<u64> {
4375        let mut json_row = Vec::with_capacity(row.len() * 2);
4376        for (id, value) in &row {
4377            json_row.push(serde_json::json!(id));
4378            json_row.push(value_to_json(value)?);
4379        }
4380        let response = self.write_response(
4381            self.client
4382                .post(self.url_segments(&["tables", table, "put"])?)
4383                .json(&serde_json::json!({ "row": json_row }))
4384                .send(),
4385            "put",
4386        )?;
4387        let response: RowIdResponse = decode_write_json(response, "put")?;
4388        parse_required_u64_text("row_id", &response.row_id)
4389            .map_err(|error| write_outcome_unknown("put", error))
4390    }
4391
4392    pub fn commit(&self, table: &str) -> ClientResult<u64> {
4393        let response = self.write_response(
4394            self.client
4395                .post(self.url_segments(&["tables", table, "commit"])?)
4396                .send(),
4397            "commit",
4398        )?;
4399        let response: EpochResponse = decode_write_json(response, "commit")?;
4400        exact_required_u64("epoch", response.epoch, &response.epoch_text)
4401            .map_err(|error| write_outcome_unknown("commit", error))
4402    }
4403
4404    // ── SQL read ──
4405
4406    pub fn sql(&self, sql: &str) -> ClientResult<Vec<RecordBatch>> {
4407        self.sql_with_options(sql, SqlClientOptions::default())
4408    }
4409
4410    pub fn sql_with_options(
4411        &self,
4412        sql: &str,
4413        mut options: SqlClientOptions,
4414    ) -> ClientResult<Vec<RecordBatch>> {
4415        if options.query_id.is_some() || options.timeout.is_some() {
4416            self.sql_cancellation_capabilities()?;
4417        }
4418        let query_id = match options.query_id.take() {
4419            Some(query_id) => query_id,
4420            None => mongreldb_query::QueryId::random()
4421                .map_err(|error| ClientError::Transport(error.to_string()))?,
4422        };
4423        let timeout_ms = options
4424            .timeout
4425            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
4426        let response = self
4427            .client
4428            .post(self.url("/sql"))
4429            .json(&SqlReq {
4430                sql: sql.to_string(),
4431                format: Some("arrow"), // Rust client decodes Arrow IPC directly
4432                query_id,
4433                timeout_ms,
4434                max_output_rows: None,
4435                max_output_bytes: None,
4436                idempotency_key: None,
4437                pagination: None,
4438            })
4439            .send();
4440        let response = match response {
4441            Ok(response) => response,
4442            Err(error) => {
4443                return Err(self.recover_after_transport_loss(query_id, error.to_string()));
4444            }
4445        };
4446        let resp = self.check_sql_response(response, query_id)?;
4447        if let Err(error) = validate_sql_query_id_header(resp.headers(), query_id) {
4448            return Err(self.recover_after_transport_loss(query_id, error));
4449        }
4450        let bytes = bounded_blocking_bytes(resp, MAX_SQL_RESPONSE_BYTES)
4451            .map_err(|error| self.recover_after_transport_loss(query_id, error.to_string()))?;
4452        read_arrow_ipc(&bytes)
4453            .map_err(|error| self.recover_after_transport_loss(query_id, error.to_string()))
4454    }
4455
4456    pub fn sql_write_idempotent(
4457        &self,
4458        sql: &str,
4459        idempotency_key: impl Into<String>,
4460    ) -> ClientResult<RemoteSqlReceipt> {
4461        self.sql_write_idempotent_with_options(sql, idempotency_key, SqlClientOptions::default())
4462    }
4463
4464    pub fn sql_write_idempotent_with_options(
4465        &self,
4466        sql: &str,
4467        idempotency_key: impl Into<String>,
4468        mut options: SqlClientOptions,
4469    ) -> ClientResult<RemoteSqlReceipt> {
4470        let idempotency_key = idempotency_key.into();
4471        if idempotency_key.is_empty() || idempotency_key.len() > 256 {
4472            return Err(ClientError::Decode(
4473                "SQL idempotency key must contain 1 to 256 bytes".into(),
4474            ));
4475        }
4476        self.sql_idempotency_capabilities()?;
4477        let query_id = match options.query_id.take() {
4478            Some(query_id) => query_id,
4479            None => mongreldb_query::QueryId::random()
4480                .map_err(|error| ClientError::Transport(error.to_string()))?,
4481        };
4482        let result =
4483            self.sql_write_idempotent_once(sql, &idempotency_key, &options, query_id, None);
4484        if result.as_ref().is_err_and(|error| error.replay) {
4485            self.sql_idempotency_capabilities()?;
4486            return self
4487                .sql_write_idempotent_once(
4488                    sql,
4489                    &idempotency_key,
4490                    &options,
4491                    fresh_query_id(query_id)?,
4492                    Some(query_id),
4493                )
4494                .map_err(|error| error.error);
4495        }
4496        result.map_err(|error| error.error)
4497    }
4498
4499    fn sql_write_idempotent_once(
4500        &self,
4501        sql: &str,
4502        idempotency_key: &str,
4503        options: &SqlClientOptions,
4504        query_id: mongreldb_query::QueryId,
4505        expected_original_query_id: Option<mongreldb_query::QueryId>,
4506    ) -> Result<RemoteSqlReceipt, IdempotentAttemptError> {
4507        let timeout_ms = options
4508            .timeout
4509            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
4510        let response = self
4511            .client
4512            .post(self.url("/sql"))
4513            .json(&SqlReq {
4514                sql: sql.to_owned(),
4515                format: None,
4516                query_id,
4517                timeout_ms,
4518                max_output_rows: None,
4519                max_output_bytes: None,
4520                idempotency_key: Some(idempotency_key.to_owned()),
4521                pagination: None,
4522            })
4523            .send();
4524        let response = match response {
4525            Ok(response) => response,
4526            Err(error) => {
4527                return Err(self.idempotent_attempt_loss(query_id, error.to_string()));
4528            }
4529        };
4530        let response = self
4531            .check_sql_response(response, query_id)
4532            .map_err(IdempotentAttemptError::final_error)?;
4533        let header_error = validate_sql_query_id_header(response.headers(), query_id).err();
4534        let bytes = bounded_blocking_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
4535            .map_err(|error| self.idempotent_attempt_loss(query_id, error))?;
4536        match decode_remote_sql_receipt(&bytes, query_id, expected_original_query_id) {
4537            Ok(receipt) if header_error.is_none() => Ok(receipt),
4538            Ok(_) => {
4539                let proof = strict_json_value(&bytes).ok().and_then(|value| {
4540                    sql_receipt_commit_proof(&value, query_id, expected_original_query_id)
4541                });
4542                match proof {
4543                    Some(proof) => Err(IdempotentAttemptError::final_error(
4544                        committed_sql_receipt_decode_error(
4545                            query_id,
4546                            proof,
4547                            header_error
4548                                .clone()
4549                                .unwrap_or_else(|| "SQL response identity is invalid".into()),
4550                        ),
4551                    )),
4552                    None => Err(self.idempotent_attempt_loss(
4553                        query_id,
4554                        header_error.unwrap_or_else(|| "SQL response identity is invalid".into()),
4555                    )),
4556                }
4557            }
4558            Err(SqlReceiptDecodeError::KnownCommit(error)) => {
4559                Err(IdempotentAttemptError::final_error(error))
4560            }
4561            Err(SqlReceiptDecodeError::Unknown(error)) => {
4562                Err(self.idempotent_attempt_loss(query_id, error))
4563            }
4564        }
4565    }
4566
4567    fn idempotent_attempt_loss(
4568        &self,
4569        query_id: mongreldb_query::QueryId,
4570        message: String,
4571    ) -> IdempotentAttemptError {
4572        let missing = self
4573            .client
4574            .get(self.url(&format!("/queries/{query_id}")))
4575            .timeout(SQL_RECOVERY_REQUEST_TIMEOUT)
4576            .send()
4577            .ok()
4578            .filter(|response| response.status() == reqwest::StatusCode::NOT_FOUND)
4579            .and_then(|response| bounded_blocking_bytes(response, MAX_CONTROL_RESPONSE_BYTES).ok())
4580            .is_some_and(|body| is_exact_query_not_found_response(&body, query_id));
4581        if missing {
4582            return IdempotentAttemptError {
4583                error: ClientError::QueryOutcomeUnknown {
4584                    query_id: query_id.to_string(),
4585                    message,
4586                    status: None,
4587                    cancel_outcome: None,
4588                },
4589                replay: true,
4590            };
4591        }
4592        IdempotentAttemptError::final_error(self.recover_after_transport_loss(query_id, message))
4593    }
4594
4595    pub fn sql_page(&self, sql: &str, mut options: SqlPageOptions) -> ClientResult<RemoteSqlPage> {
4596        validate_sql_page_options(&options)?;
4597        self.sql_pagination_capabilities()?;
4598        let query_id = match options.query_id.take() {
4599            Some(query_id) => query_id,
4600            None => mongreldb_query::QueryId::random()
4601                .map_err(|error| ClientError::Transport(error.to_string()))?,
4602        };
4603        let timeout_ms = options
4604            .timeout
4605            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
4606        let response = self
4607            .client
4608            .post(self.url("/sql"))
4609            .json(&SqlReq {
4610                sql: sql.to_owned(),
4611                format: None,
4612                query_id,
4613                timeout_ms,
4614                max_output_rows: options.max_output_rows,
4615                max_output_bytes: options.max_output_bytes,
4616                idempotency_key: None,
4617                pagination: Some(SqlPaginationReq {
4618                    page_size_rows: options.page_size_rows,
4619                    projection: options.projection.clone(),
4620                    max_page_bytes: options.max_page_bytes,
4621                    max_page_tokens: options.max_page_tokens,
4622                }),
4623            })
4624            .send();
4625        let response = match response {
4626            Ok(response) => response,
4627            Err(error) => {
4628                return Err(self.recover_after_transport_loss(query_id, error.to_string()));
4629            }
4630        };
4631        let page = bounded_blocking_bytes(
4632            {
4633                let response = self.check_sql_response(response, query_id)?;
4634                if let Err(error) = validate_sql_query_id_header(response.headers(), query_id) {
4635                    return Err(self.recover_after_transport_loss(query_id, error));
4636                }
4637                response
4638            },
4639            MAX_SQL_RESPONSE_BYTES,
4640        )
4641        .and_then(|bytes| {
4642            strict_json::<RemoteSqlPage>(&bytes, "SQL page").map_err(|error| error.to_string())
4643        })
4644        .and_then(|page| validate_remote_sql_page(page, Some(&options)));
4645        page.map_err(|error| self.recover_after_transport_loss(query_id, error))
4646    }
4647
4648    pub fn continue_sql_page(
4649        &self,
4650        cursor: &str,
4651        mut options: RemoteSqlControlOptions,
4652    ) -> ClientResult<RemoteSqlPage> {
4653        if cursor.is_empty() || cursor.len() > 2_048 {
4654            return Err(ClientError::Decode(
4655                "SQL continuation cursor must contain 1 to 2048 bytes".into(),
4656            ));
4657        }
4658        self.sql_pagination_capabilities()?;
4659        let query_id = match options.query_id.take() {
4660            Some(query_id) => query_id,
4661            None => mongreldb_query::QueryId::random()
4662                .map_err(|error| ClientError::Transport(error.to_string()))?,
4663        };
4664        if options.timeout == Some(std::time::Duration::ZERO) {
4665            return Err(ClientError::Decode("timeout must be positive".into()));
4666        }
4667        let timeout_ms = options
4668            .timeout
4669            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
4670        let response = match self
4671            .client
4672            .post(self.url("/sql/continue"))
4673            .json(&serde_json::json!({
4674                "cursor": cursor,
4675                "operation_id": query_id,
4676                "timeout_ms": timeout_ms,
4677            }))
4678            .send()
4679        {
4680            Ok(response) => response,
4681            Err(error) => {
4682                return Err(self.recover_after_transport_loss(query_id, error.to_string()))
4683            }
4684        };
4685        let response = self.check_sql_response(response, query_id)?;
4686        validate_sql_query_id_header(response.headers(), query_id)
4687            .map_err(|error| client_serialization_error(Some(query_id), error))?;
4688        bounded_blocking_bytes(response, MAX_SQL_RESPONSE_BYTES)
4689            .and_then(|bytes| {
4690                strict_json::<RemoteSqlPage>(&bytes, "SQL page").map_err(|error| error.to_string())
4691            })
4692            .and_then(|page| validate_remote_sql_page(page, None))
4693            .map_err(|error| client_serialization_error(Some(query_id), error))
4694    }
4695
4696    pub fn start_sql(
4697        &self,
4698        sql: impl Into<String>,
4699        mut options: SqlClientOptions,
4700    ) -> ClientResult<RemoteSqlQueryHandle> {
4701        self.sql_cancellation_capabilities()?;
4702        let query_id = match options.query_id {
4703            Some(query_id) => query_id,
4704            None => mongreldb_query::QueryId::random()
4705                .map_err(|error| ClientError::Transport(error.to_string()))?,
4706        };
4707        options.query_id = Some(query_id);
4708        let client = self.clone();
4709        let cancel_client = self.clone();
4710        let sql = sql.into();
4711        let result = std::thread::Builder::new()
4712            .name(format!("mongreldb-sql-{query_id}"))
4713            .spawn(move || client.sql_with_options(&sql, options))
4714            .map_err(|error| ClientError::Transport(error.to_string()))?;
4715        Ok(RemoteSqlQueryHandle {
4716            query_id,
4717            client: cancel_client,
4718            result,
4719        })
4720    }
4721
4722    pub fn cancel_sql(
4723        &self,
4724        query_id: mongreldb_query::QueryId,
4725    ) -> ClientResult<RemoteCancelOutcome> {
4726        self.sql_cancellation_capabilities()?;
4727        let response = self
4728            .client
4729            .post(self.url(&format!("/queries/{query_id}/cancel")))
4730            .send()?;
4731        let status = response.status();
4732        if !matches!(
4733            status,
4734            reqwest::StatusCode::OK
4735                | reqwest::StatusCode::ACCEPTED
4736                | reqwest::StatusCode::CONFLICT
4737                | reqwest::StatusCode::NOT_FOUND
4738        ) {
4739            return match self.check(response) {
4740                Err(error) => Err(error),
4741                Ok(_) => Err(ClientError::Decode(
4742                    "unexpected successful cancellation response".into(),
4743                )),
4744            };
4745        }
4746        let body: serde_json::Value =
4747            decode_blocking_json(response, MAX_CONTROL_RESPONSE_BYTES, "cancellation")?;
4748        decode_cancel_outcome(&body, query_id, status.as_u16())
4749    }
4750
4751    pub fn query_status(
4752        &self,
4753        query_id: mongreldb_query::QueryId,
4754    ) -> ClientResult<RemoteQueryStatus> {
4755        self.sql_cancellation_capabilities()?;
4756        let response = self
4757            .client
4758            .get(self.url(&format!("/queries/{query_id}")))
4759            .send()?;
4760        let status = decode_blocking_json(
4761            self.check_query_status_response(response, query_id)?,
4762            MAX_CONTROL_RESPONSE_BYTES,
4763            "query status",
4764        )?;
4765        validate_remote_query_status(status, query_id).map_err(ClientError::Decode)
4766    }
4767
4768    fn query_status_optional(
4769        &self,
4770        query_id: mongreldb_query::QueryId,
4771        timeout: std::time::Duration,
4772    ) -> Option<RemoteQueryStatus> {
4773        let response = self
4774            .client
4775            .get(self.url(&format!("/queries/{query_id}")))
4776            .timeout(timeout)
4777            .send()
4778            .ok()?;
4779        if !response.status().is_success() {
4780            return None;
4781        }
4782        let status =
4783            decode_blocking_json(response, MAX_CONTROL_RESPONSE_BYTES, "query status").ok()?;
4784        validate_remote_query_status(status, query_id).ok()
4785    }
4786
4787    fn cancel_sql_for_recovery(
4788        &self,
4789        query_id: mongreldb_query::QueryId,
4790        timeout: std::time::Duration,
4791    ) -> Option<RemoteCancelOutcome> {
4792        let response = self
4793            .client
4794            .post(self.url(&format!("/queries/{query_id}/cancel")))
4795            .timeout(timeout)
4796            .send()
4797            .ok()?;
4798        let status = response.status();
4799        if !matches!(
4800            status,
4801            reqwest::StatusCode::OK
4802                | reqwest::StatusCode::ACCEPTED
4803                | reqwest::StatusCode::CONFLICT
4804                | reqwest::StatusCode::NOT_FOUND
4805        ) {
4806            return None;
4807        }
4808        let body =
4809            decode_blocking_json(response, MAX_CONTROL_RESPONSE_BYTES, "cancellation").ok()?;
4810        decode_cancel_outcome(&body, query_id, status.as_u16()).ok()
4811    }
4812
4813    fn recover_after_transport_loss(
4814        &self,
4815        query_id: mongreldb_query::QueryId,
4816        message: String,
4817    ) -> ClientError {
4818        let deadline = std::time::Instant::now() + SQL_RECOVERY_WINDOW;
4819        let mut status = self.query_status_optional(query_id, SQL_RECOVERY_REQUEST_TIMEOUT);
4820        if let Some(decisive) = status
4821            .as_ref()
4822            .filter(|status| recovery_status_is_decisive(status))
4823        {
4824            return recovered_query_error(decisive.clone(), message);
4825        }
4826        let cancel_outcome = self.cancel_sql_for_recovery(
4827            query_id,
4828            deadline
4829                .saturating_duration_since(std::time::Instant::now())
4830                .min(SQL_RECOVERY_REQUEST_TIMEOUT),
4831        );
4832        if status.is_none() && cancel_outcome == Some(RemoteCancelOutcome::NotFound) {
4833            return ClientError::QueryOutcomeUnknown {
4834                query_id: query_id.to_string(),
4835                message,
4836                status: None,
4837                cancel_outcome,
4838            };
4839        }
4840        while !status.as_ref().is_some_and(recovery_status_is_decisive)
4841            && std::time::Instant::now() < deadline
4842        {
4843            std::thread::sleep(
4844                deadline
4845                    .saturating_duration_since(std::time::Instant::now())
4846                    .min(SQL_RECOVERY_POLL_INTERVAL),
4847            );
4848            let remaining = deadline.saturating_duration_since(std::time::Instant::now());
4849            if remaining.is_zero() {
4850                break;
4851            }
4852            status = self
4853                .query_status_optional(query_id, remaining.min(SQL_RECOVERY_REQUEST_TIMEOUT))
4854                .or(status);
4855        }
4856        if let Some(decisive) = status
4857            .as_ref()
4858            .filter(|status| recovery_status_is_decisive(status))
4859        {
4860            return recovered_query_error(decisive.clone(), message);
4861        }
4862        ClientError::QueryOutcomeUnknown {
4863            query_id: query_id.to_string(),
4864            message,
4865            status: status.map(Box::new),
4866            cancel_outcome,
4867        }
4868    }
4869
4870    // ── Atomic txn (legacy, raw) ──
4871
4872    pub fn txn(&self, ops: Vec<TxnOp>) -> ClientResult<()> {
4873        let response = self.write_response(
4874            self.client
4875                .post(self.url("/txn"))
4876                .json(&serde_json::json!({ "ops": ops }))
4877                .send(),
4878            "transaction",
4879        )?;
4880        let response: CommittedWriteResponse = decode_write_json(response, "transaction")?;
4881        validate_committed_write(&response, "transaction")?;
4882        Ok(())
4883    }
4884
4885    // ── Typed Kit surface ──
4886
4887    /// Fetch one table's schema + constraint metadata (`GET /kit/schema/{t}`).
4888    pub fn kit_schema(&self, table: &str) -> ClientResult<TableSchemaInfo> {
4889        let resp = self
4890            .client
4891            .get(self.url_segments(&["kit", "schema", table])?)
4892            .send()?;
4893        decode_blocking_json(self.check(resp)?, MAX_CONTROL_RESPONSE_BYTES, "Kit schema")
4894    }
4895
4896    /// Run a typed atomic batch (`POST /kit/txn`). Constraint violations and
4897    /// conflicts return [`ClientError::Kit`] with the matching code.
4898    pub fn kit_txn(&self, req: &KitTxnRequest) -> ClientResult<KitTxnResponse> {
4899        let resp = self
4900            .client
4901            .post(self.url("/kit/txn"))
4902            .json(req)
4903            .send()
4904            .map_err(|error| {
4905                kit_txn_outcome_unknown(format!(
4906                    "Kit transaction transport failed before outcome confirmation: {error}"
4907                ))
4908            })?;
4909        if !resp.status().is_success() {
4910            let status = resp.status().as_u16();
4911            if matches!(status, 401 | 403) {
4912                return Err(kit_txn_auth_error(status));
4913            }
4914            let body = bounded_blocking_bytes(resp, MAX_SQL_RESPONSE_BYTES).map_err(|error| {
4915                kit_txn_outcome_unknown(format!(
4916                    "Kit transaction error response could not be read: {error}"
4917                ))
4918            })?;
4919            return Err(decode_kit_txn_http_error(status, &body, req));
4920        }
4921        let status = resp.status().as_u16();
4922        let body = bounded_blocking_bytes(resp, MAX_SQL_RESPONSE_BYTES).map_err(|error| {
4923            kit_txn_outcome_unknown(format!(
4924                "Kit transaction success response could not be read: {error}"
4925            ))
4926        })?;
4927        decode_kit_txn_success(&body, req, status)
4928    }
4929
4930    /// Run a native typed query (`POST /kit/query`) returning physical row ids
4931    /// and typed cells. Conditions intersect in the row-id space; this is the
4932    /// native counterpart to SQL reads (which hide row ids).
4933    pub fn kit_query(&self, req: &KitQueryRequest) -> ClientResult<KitQueryResponse> {
4934        let resp = self.client.post(self.url("/kit/query")).json(req).send()?;
4935        let response =
4936            decode_blocking_json(self.check(resp)?, MAX_SQL_RESPONSE_BYTES, "Kit query")?;
4937        validate_kit_query_response(response, req)
4938    }
4939
4940    pub fn kit_retrieve(&self, req: &KitRetrieveRequest) -> ClientResult<KitRetrieveResponse> {
4941        self.kit_retrieve_with_options(req, &AiExecutionOptions::default())
4942    }
4943
4944    pub fn kit_retrieve_with_options(
4945        &self,
4946        req: &KitRetrieveRequest,
4947        options: &AiExecutionOptions,
4948    ) -> ClientResult<KitRetrieveResponse> {
4949        let resp = self
4950            .client
4951            .post(self.url("/kit/retrieve"))
4952            .json(&WithAiExecutionOptions {
4953                request: req,
4954                options,
4955            })
4956            .send()?;
4957        decode_blocking_json(self.check(resp)?, MAX_SQL_RESPONSE_BYTES, "Kit retrieve")
4958    }
4959
4960    pub fn kit_ann_rerank(&self, req: &KitAnnRerankRequest) -> ClientResult<KitAnnRerankResponse> {
4961        self.kit_ann_rerank_with_options(req, &AiExecutionOptions::default())
4962    }
4963
4964    pub fn kit_ann_rerank_with_options(
4965        &self,
4966        req: &KitAnnRerankRequest,
4967        options: &AiExecutionOptions,
4968    ) -> ClientResult<KitAnnRerankResponse> {
4969        let resp = self
4970            .client
4971            .post(self.url("/kit/ann_rerank"))
4972            .json(&WithAiExecutionOptions {
4973                request: req,
4974                options,
4975            })
4976            .send()?;
4977        decode_blocking_json(self.check(resp)?, MAX_SQL_RESPONSE_BYTES, "Kit ANN rerank")
4978    }
4979
4980    pub fn kit_set_similarity(
4981        &self,
4982        req: &KitSetSimilarityRequest,
4983    ) -> ClientResult<KitSetSimilarityResponse> {
4984        self.kit_set_similarity_with_options(req, &AiExecutionOptions::default())
4985    }
4986
4987    pub fn kit_set_similarity_with_options(
4988        &self,
4989        req: &KitSetSimilarityRequest,
4990        options: &AiExecutionOptions,
4991    ) -> ClientResult<KitSetSimilarityResponse> {
4992        let resp = self
4993            .client
4994            .post(self.url("/kit/set_similarity"))
4995            .json(&WithAiExecutionOptions {
4996                request: req,
4997                options,
4998            })
4999            .send()?;
5000        decode_blocking_json(
5001            self.check(resp)?,
5002            MAX_SQL_RESPONSE_BYTES,
5003            "Kit set similarity",
5004        )
5005    }
5006
5007    pub fn kit_search(&self, req: &KitSearchRequest) -> ClientResult<KitSearchResponse> {
5008        let resp = self.client.post(self.url("/kit/search")).json(req).send()?;
5009        decode_blocking_json(self.check(resp)?, MAX_SQL_RESPONSE_BYTES, "Kit search")
5010    }
5011
5012    pub fn kit_ai_metrics(&self) -> ClientResult<serde_json::Value> {
5013        let resp = self.client.get(self.url("/kit/ai/metrics")).send()?;
5014        decode_blocking_json(
5015            self.check(resp)?,
5016            MAX_CONTROL_RESPONSE_BYTES,
5017            "Kit AI metrics",
5018        )
5019    }
5020
5021    /// Create a constraint-bearing table over HTTP (`POST /kit/create_table`).
5022    /// `body` is the full request JSON — `{name, columns:[{id,name,ty,
5023    /// primary_key,nullable,auto_increment,…}], constraints:{uniques,…,
5024    /// foreign_keys,…, checks:[{id,name,expr}]}}`. Returns the assigned table id.
5025    pub fn kit_create_table(&self, body: &serde_json::Value) -> ClientResult<u64> {
5026        let response = self.write_response(
5027            self.client
5028                .post(self.url("/kit/create_table"))
5029                .json(body)
5030                .send(),
5031            "Kit create table",
5032        )?;
5033        let response: TableIdResponse = decode_write_json(response, "Kit create table")?;
5034        exact_required_u64("table_id", response.table_id, &response.table_id_text)
5035            .map_err(|error| write_outcome_unknown("Kit create table", error))
5036    }
5037
5038    pub fn procedures(&self) -> ClientResult<Vec<mongreldb_core::StoredProcedure>> {
5039        let resp = self.client.get(self.url("/procedures")).send()?;
5040        let response: ProceduresResponse = decode_blocking_json(
5041            self.check(resp)?,
5042            MAX_CONTROL_RESPONSE_BYTES,
5043            "procedure list",
5044        )?;
5045        Ok(response.procedures)
5046    }
5047
5048    pub fn procedure(&self, name: &str) -> ClientResult<mongreldb_core::StoredProcedure> {
5049        let resp = self
5050            .client
5051            .get(self.url_segments(&["procedures", name])?)
5052            .send()?;
5053        let response: ProcedureResponse =
5054            decode_blocking_json(self.check(resp)?, MAX_CONTROL_RESPONSE_BYTES, "procedure")?;
5055        Ok(response.procedure)
5056    }
5057
5058    pub fn create_procedure(
5059        &self,
5060        procedure: mongreldb_core::StoredProcedure,
5061    ) -> ClientResult<mongreldb_core::StoredProcedure> {
5062        let expected_name = procedure.name.clone();
5063        let expected_definition = procedure.clone();
5064        let response = self.write_response(
5065            self.client
5066                .post(self.url("/procedures"))
5067                .json(&ProcedureRequest { procedure })
5068                .send(),
5069            "create procedure",
5070        )?;
5071        let response: ProcedureResponse = decode_write_json(response, "create procedure")?;
5072        if response.status != "ok"
5073            || response.procedure.name != expected_name
5074            || response.procedure.version == 0
5075            || response.procedure.created_epoch == 0
5076            || response.procedure.updated_epoch < response.procedure.created_epoch
5077            || response.procedure.checksum.is_empty()
5078            || response.procedure.validate().is_err()
5079            || !procedure_definition_matches(&response.procedure, &expected_definition)
5080        {
5081            return Err(write_outcome_unknown(
5082                "create procedure",
5083                "success response does not match the requested procedure",
5084            ));
5085        }
5086        Ok(response.procedure)
5087    }
5088
5089    pub fn replace_procedure(
5090        &self,
5091        name: &str,
5092        procedure: mongreldb_core::StoredProcedure,
5093    ) -> ClientResult<mongreldb_core::StoredProcedure> {
5094        let expected_definition = procedure.clone();
5095        let response = self.write_response(
5096            self.client
5097                .put(self.url_segments(&["procedures", name])?)
5098                .json(&ProcedureRequest { procedure })
5099                .send(),
5100            "replace procedure",
5101        )?;
5102        let response: ProcedureResponse = decode_write_json(response, "replace procedure")?;
5103        if response.status != "ok"
5104            || response.procedure.name != name
5105            || response.procedure.version == 0
5106            || response.procedure.created_epoch == 0
5107            || response.procedure.updated_epoch < response.procedure.created_epoch
5108            || response.procedure.checksum.is_empty()
5109            || response.procedure.validate().is_err()
5110            || !procedure_definition_matches(&response.procedure, &expected_definition)
5111        {
5112            return Err(write_outcome_unknown(
5113                "replace procedure",
5114                "success response does not match the requested procedure",
5115            ));
5116        }
5117        Ok(response.procedure)
5118    }
5119
5120    pub fn drop_procedure(&self, name: &str) -> ClientResult<()> {
5121        let response = self.write_response(
5122            self.client
5123                .delete(self.url_segments(&["procedures", name])?)
5124                .send(),
5125            "drop procedure",
5126        )?;
5127        let response: CommittedWriteResponse = decode_write_json(response, "drop procedure")?;
5128        validate_committed_write(&response, "drop procedure")?;
5129        Ok(())
5130    }
5131
5132    pub fn call_procedure(
5133        &self,
5134        name: &str,
5135        req: &ProcedureCallRequest,
5136    ) -> ClientResult<ProcedureCallResponse> {
5137        let response = self.write_response(
5138            self.client
5139                .post(self.url_segments(&["procedures", name, "call"])?)
5140                .json(req)
5141                .send(),
5142            "procedure call",
5143        )?;
5144        let response = decode_write_json(response, "procedure call")?;
5145        validate_procedure_call_response(response)
5146            .map_err(|error| write_outcome_unknown("procedure call", error))
5147    }
5148
5149    pub fn kit_call_procedure(
5150        &self,
5151        name: &str,
5152        req: &ProcedureCallRequest,
5153    ) -> ClientResult<ProcedureCallResponse> {
5154        let response = self.write_response(
5155            self.client
5156                .post(self.url_segments(&["kit", "procedures", name, "call"])?)
5157                .json(req)
5158                .send(),
5159            "Kit procedure call",
5160        )?;
5161        let response = decode_write_json(response, "Kit procedure call")?;
5162        validate_procedure_call_response(response)
5163            .map_err(|error| write_outcome_unknown("Kit procedure call", error))
5164    }
5165
5166    pub fn triggers(&self) -> ClientResult<Vec<mongreldb_core::StoredTrigger>> {
5167        let resp = self.client.get(self.url("/triggers")).send()?;
5168        let response: TriggersResponse = decode_blocking_json(
5169            self.check(resp)?,
5170            MAX_CONTROL_RESPONSE_BYTES,
5171            "trigger list",
5172        )?;
5173        Ok(response.triggers)
5174    }
5175
5176    pub fn trigger(&self, name: &str) -> ClientResult<mongreldb_core::StoredTrigger> {
5177        let resp = self
5178            .client
5179            .get(self.url_segments(&["triggers", name])?)
5180            .send()?;
5181        let response: TriggerResponse =
5182            decode_blocking_json(self.check(resp)?, MAX_CONTROL_RESPONSE_BYTES, "trigger")?;
5183        Ok(response.trigger)
5184    }
5185
5186    pub fn create_trigger(
5187        &self,
5188        trigger: mongreldb_core::StoredTrigger,
5189    ) -> ClientResult<mongreldb_core::StoredTrigger> {
5190        self.create_trigger_with_idempotency_key(trigger, None::<String>)
5191    }
5192
5193    pub fn create_trigger_with_idempotency_key(
5194        &self,
5195        trigger: mongreldb_core::StoredTrigger,
5196        idempotency_key: Option<impl Into<String>>,
5197    ) -> ClientResult<mongreldb_core::StoredTrigger> {
5198        let expected_name = trigger.name.clone();
5199        let expected_definition = trigger.clone();
5200        let response = self.write_response(
5201            self.client
5202                .post(self.url("/triggers"))
5203                .json(&TriggerRequest {
5204                    trigger,
5205                    idempotency_key: idempotency_key.map(Into::into),
5206                })
5207                .send(),
5208            "create trigger",
5209        )?;
5210        let response: TriggerResponse = decode_write_json(response, "create trigger")?;
5211        if response.status.as_deref() != Some("ok")
5212            || response.trigger.name != expected_name
5213            || response.trigger.version == 0
5214            || response.trigger.created_epoch == 0
5215            || response.trigger.updated_epoch < response.trigger.created_epoch
5216            || response.trigger.checksum.is_empty()
5217            || response.trigger.validate().is_err()
5218            || !trigger_definition_matches(&response.trigger, &expected_definition)
5219        {
5220            return Err(write_outcome_unknown(
5221                "create trigger",
5222                "success response does not match the requested trigger",
5223            ));
5224        }
5225        Ok(response.trigger)
5226    }
5227
5228    pub fn replace_trigger(
5229        &self,
5230        name: &str,
5231        trigger: mongreldb_core::StoredTrigger,
5232    ) -> ClientResult<mongreldb_core::StoredTrigger> {
5233        self.replace_trigger_with_idempotency_key(name, trigger, None::<String>)
5234    }
5235
5236    pub fn replace_trigger_with_idempotency_key(
5237        &self,
5238        name: &str,
5239        trigger: mongreldb_core::StoredTrigger,
5240        idempotency_key: Option<impl Into<String>>,
5241    ) -> ClientResult<mongreldb_core::StoredTrigger> {
5242        let expected_definition = trigger.clone();
5243        let response = self.write_response(
5244            self.client
5245                .put(self.url_segments(&["triggers", name])?)
5246                .json(&TriggerRequest {
5247                    trigger,
5248                    idempotency_key: idempotency_key.map(Into::into),
5249                })
5250                .send(),
5251            "replace trigger",
5252        )?;
5253        let response: TriggerResponse = decode_write_json(response, "replace trigger")?;
5254        if response.status.as_deref() != Some("ok")
5255            || response.trigger.name != name
5256            || response.trigger.version == 0
5257            || response.trigger.created_epoch == 0
5258            || response.trigger.updated_epoch < response.trigger.created_epoch
5259            || response.trigger.checksum.is_empty()
5260            || response.trigger.validate().is_err()
5261            || !trigger_definition_matches(&response.trigger, &expected_definition)
5262        {
5263            return Err(write_outcome_unknown(
5264                "replace trigger",
5265                "success response does not match the requested trigger",
5266            ));
5267        }
5268        Ok(response.trigger)
5269    }
5270
5271    pub fn drop_trigger(&self, name: &str) -> ClientResult<()> {
5272        self.drop_trigger_with_idempotency_key(name, None::<String>)
5273    }
5274
5275    pub fn drop_trigger_with_idempotency_key(
5276        &self,
5277        name: &str,
5278        idempotency_key: Option<impl Into<String>>,
5279    ) -> ClientResult<()> {
5280        let mut request = self.client.delete(self.url_segments(&["triggers", name])?);
5281        if let Some(idempotency_key) = idempotency_key {
5282            request = request.header("Idempotency-Key", idempotency_key.into());
5283        }
5284        let response = self.write_response(request.send(), "drop trigger")?;
5285        let response: TriggerDropResponse = decode_write_json(response, "drop trigger")?;
5286        validate_trigger_drop_response(&response, name)?;
5287        Ok(())
5288    }
5289}
5290
5291impl AsyncMongrelClient {
5292    pub fn builder(url: impl AsRef<str>) -> AsyncMongrelClientBuilder {
5293        let base_url = sanitized_base_url(url.as_ref());
5294        AsyncMongrelClientBuilder {
5295            invalid_base_url: base_url.is_none(),
5296            base_url: base_url.unwrap_or_default(),
5297            authorization: None,
5298            invalid_authorization: false,
5299            connect_timeout: None,
5300            request_timeout: None,
5301            pool_idle_timeout: None,
5302        }
5303    }
5304
5305    pub fn new(url: &str) -> ClientResult<Self> {
5306        Self::builder(url).build()
5307    }
5308
5309    pub fn with_options(url: impl AsRef<str>, options: RemoteOptions) -> ClientResult<Self> {
5310        let mut builder = Self::builder(url);
5311        if let Some(timeout) = options.transport_timeout {
5312            builder = builder.request_timeout(timeout);
5313        }
5314        builder = match options.auth {
5315            Some(RemoteAuth::Bearer(token)) => builder.bearer_token(token.expose_secret()),
5316            Some(RemoteAuth::Basic { username, password }) => {
5317                builder.basic_auth(username, password.expose_secret())
5318            }
5319            None => builder,
5320        };
5321        builder.build()
5322    }
5323
5324    pub fn try_with_bearer_token(mut self, token: impl AsRef<str>) -> ClientResult<Self> {
5325        self.client = async_http_client(Some(bearer_header(token.as_ref())?), self.transport)?;
5326        Ok(self)
5327    }
5328
5329    pub fn with_bearer_token(self, token: impl AsRef<str>) -> ClientResult<Self> {
5330        self.try_with_bearer_token(token)
5331    }
5332
5333    pub fn try_with_basic_auth(
5334        mut self,
5335        username: impl AsRef<str>,
5336        password: impl AsRef<str>,
5337    ) -> ClientResult<Self> {
5338        self.client = async_http_client(
5339            Some(basic_header(username.as_ref(), password.as_ref())?),
5340            self.transport,
5341        )?;
5342        Ok(self)
5343    }
5344
5345    pub fn with_basic_auth(
5346        self,
5347        username: impl AsRef<str>,
5348        password: impl AsRef<str>,
5349    ) -> ClientResult<Self> {
5350        self.try_with_basic_auth(username, password)
5351    }
5352
5353    fn url(&self, path: &str) -> String {
5354        format!("{}{path}", self.base_url)
5355    }
5356
5357    async fn check(&self, response: reqwest::Response) -> ClientResult<reqwest::Response> {
5358        let status = response.status();
5359        if status.is_success() {
5360            return Ok(response);
5361        }
5362        let status_u16 = status.as_u16();
5363        let body = bounded_async_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
5364            .await
5365            .map_err(|error| {
5366                ClientError::Decode(format!("invalid HTTP error response: {error}"))
5367            })?;
5368        Err(decode_http_error(status_u16, &body))
5369    }
5370
5371    async fn check_sql_response(
5372        &self,
5373        response: reqwest::Response,
5374        query_id: mongreldb_query::QueryId,
5375    ) -> ClientResult<reqwest::Response> {
5376        let status = response.status();
5377        let pre_handler_auth = matches!(
5378            status,
5379            reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN
5380        );
5381        if status.is_success() {
5382            return Ok(response);
5383        }
5384        let header_error = (!pre_handler_auth)
5385            .then(|| validate_sql_query_id_header(response.headers(), query_id).err())
5386            .flatten();
5387        let status = status.as_u16();
5388        let body = match bounded_async_bytes(response, MAX_CONTROL_RESPONSE_BYTES).await {
5389            Ok(body) => body,
5390            Err(error) => {
5391                if pre_handler_auth {
5392                    return Err(ClientError::Http {
5393                        status,
5394                        body: format!("unreadable authentication error response: {error}"),
5395                    });
5396                }
5397                return Err(self
5398                    .recover_after_transport_loss(query_id, error.to_string())
5399                    .await);
5400            }
5401        };
5402        if pre_handler_auth {
5403            return Err(decode_http_error(status, &body));
5404        }
5405        match decode_sql_http_error(status, &body, query_id) {
5406            Ok(error) if client_error_proves_commit(&error) => Err(error),
5407            Ok(error) if header_error.is_none() => Err(error),
5408            Ok(_) => Err(self
5409                .recover_after_transport_loss(
5410                    query_id,
5411                    header_error.unwrap_or_else(|| "SQL response identity is invalid".into()),
5412                )
5413                .await),
5414            Err(error) => Err(self.recover_after_transport_loss(query_id, error).await),
5415        }
5416    }
5417
5418    async fn check_query_status_response(
5419        &self,
5420        response: reqwest::Response,
5421        query_id: mongreldb_query::QueryId,
5422    ) -> ClientResult<reqwest::Response> {
5423        let status = response.status();
5424        if status.is_success() {
5425            return Ok(response);
5426        }
5427        let status = status.as_u16();
5428        let body = bounded_async_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
5429            .await
5430            .map_err(ClientError::Decode)?;
5431        if matches!(status, 401 | 403) {
5432            return Err(decode_http_error(status, &body));
5433        }
5434        match decode_sql_http_error(status, &body, query_id) {
5435            Ok(error) => Err(error),
5436            Err(error) => Err(ClientError::Decode(error)),
5437        }
5438    }
5439
5440    pub async fn health(&self) -> ClientResult<String> {
5441        let response = self.client.get(self.url("/health")).send().await?;
5442        let bytes = bounded_async_bytes(self.check(response).await?, MAX_CONTROL_RESPONSE_BYTES)
5443            .await
5444            .map_err(ClientError::Transport)?;
5445        String::from_utf8(bytes)
5446            .map_err(|_| ClientError::Decode("invalid health response: non-UTF-8 body".into()))
5447    }
5448
5449    pub async fn capabilities(&self) -> ClientResult<ServerCapabilities> {
5450        let response = self.client.get(self.url("/capabilities")).send().await?;
5451        decode_async_json(
5452            self.check(response).await?,
5453            MAX_CONTROL_RESPONSE_BYTES,
5454            "capabilities",
5455        )
5456        .await
5457    }
5458
5459    pub async fn sql_cancellation_capabilities(&self) -> ClientResult<SqlCancellationCapabilities> {
5460        let capabilities = self.capabilities().await?.sql_cancellation;
5461        if capabilities.version < 2
5462            || !capabilities.client_query_ids
5463            || !capabilities.cancel_endpoint
5464            || !capabilities.query_status
5465            || !capabilities.pre_registration_cancel
5466        {
5467            return Err(capability_unsupported(
5468                "server does not support SQL cancellation capability version 2",
5469            ));
5470        }
5471        Ok(capabilities)
5472    }
5473
5474    pub async fn sql_idempotency_capabilities(&self) -> ClientResult<SqlIdempotencyCapabilities> {
5475        let capabilities =
5476            self.capabilities().await?.sql_idempotency.ok_or_else(|| {
5477                capability_unsupported("server does not advertise SQL idempotency")
5478            })?;
5479        if capabilities.version < 1
5480            || !capabilities.durable_pre_execution_intent
5481            || !capabilities.replay_committed_receipt
5482            || !capabilities.indeterminate_never_reexecutes
5483        {
5484            return Err(capability_unsupported(
5485                "server does not support safe SQL idempotency capability version 1",
5486            ));
5487        }
5488        Ok(capabilities)
5489    }
5490
5491    pub async fn sql_pagination_capabilities(&self) -> ClientResult<SqlPaginationCapabilities> {
5492        let capabilities =
5493            self.capabilities().await?.sql_pagination.ok_or_else(|| {
5494                capability_unsupported("server does not advertise SQL pagination")
5495            })?;
5496        if capabilities.version < 1
5497            || capabilities.continuation_endpoint != "/sql/continue"
5498            || !capabilities.retained_snapshot
5499            || !capabilities.projection_required
5500            || !capabilities.byte_and_token_hints
5501        {
5502            return Err(capability_unsupported(
5503                "server does not support SQL pagination capability version 1",
5504            ));
5505        }
5506        Ok(capabilities)
5507    }
5508
5509    pub async fn sql(&self, sql: &str) -> ClientResult<Vec<RecordBatch>> {
5510        self.sql_with_options(sql, SqlClientOptions::default())
5511            .await
5512    }
5513
5514    pub async fn sql_with_options(
5515        &self,
5516        sql: &str,
5517        mut options: SqlClientOptions,
5518    ) -> ClientResult<Vec<RecordBatch>> {
5519        if options.query_id.is_some() || options.timeout.is_some() {
5520            self.sql_cancellation_capabilities().await?;
5521        }
5522        let query_id = match options.query_id.take() {
5523            Some(query_id) => query_id,
5524            None => mongreldb_query::QueryId::random()
5525                .map_err(|error| ClientError::Transport(error.to_string()))?,
5526        };
5527        let timeout_ms = options
5528            .timeout
5529            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
5530        let response = self
5531            .client
5532            .post(self.url("/sql"))
5533            .json(&SqlReq {
5534                sql: sql.to_string(),
5535                format: Some("arrow"),
5536                query_id,
5537                timeout_ms,
5538                max_output_rows: None,
5539                max_output_bytes: None,
5540                idempotency_key: None,
5541                pagination: None,
5542            })
5543            .send()
5544            .await;
5545        let response = match response {
5546            Ok(response) => response,
5547            Err(error) => {
5548                return Err(self
5549                    .recover_after_transport_loss(query_id, error.to_string())
5550                    .await);
5551            }
5552        };
5553        let response = self.check_sql_response(response, query_id).await?;
5554        if let Err(error) = validate_sql_query_id_header(response.headers(), query_id) {
5555            return Err(self.recover_after_transport_loss(query_id, error).await);
5556        }
5557        let bytes = match bounded_async_bytes(response, MAX_SQL_RESPONSE_BYTES).await {
5558            Ok(bytes) => bytes,
5559            Err(error) => {
5560                return Err(self
5561                    .recover_after_transport_loss(query_id, error.to_string())
5562                    .await)
5563            }
5564        };
5565        match read_arrow_ipc(&bytes) {
5566            Ok(batches) => Ok(batches),
5567            Err(error) => Err(self
5568                .recover_after_transport_loss(query_id, error.to_string())
5569                .await),
5570        }
5571    }
5572
5573    pub async fn sql_write_idempotent(
5574        &self,
5575        sql: &str,
5576        idempotency_key: impl Into<String>,
5577    ) -> ClientResult<RemoteSqlReceipt> {
5578        self.sql_write_idempotent_with_options(sql, idempotency_key, SqlClientOptions::default())
5579            .await
5580    }
5581
5582    pub async fn sql_write_idempotent_with_options(
5583        &self,
5584        sql: &str,
5585        idempotency_key: impl Into<String>,
5586        mut options: SqlClientOptions,
5587    ) -> ClientResult<RemoteSqlReceipt> {
5588        let idempotency_key = idempotency_key.into();
5589        if idempotency_key.is_empty() || idempotency_key.len() > 256 {
5590            return Err(ClientError::Decode(
5591                "SQL idempotency key must contain 1 to 256 bytes".into(),
5592            ));
5593        }
5594        self.sql_idempotency_capabilities().await?;
5595        let query_id = match options.query_id.take() {
5596            Some(query_id) => query_id,
5597            None => mongreldb_query::QueryId::random()
5598                .map_err(|error| ClientError::Transport(error.to_string()))?,
5599        };
5600        let result = self
5601            .sql_write_idempotent_once(sql, &idempotency_key, &options, query_id, None)
5602            .await;
5603        if result.as_ref().is_err_and(|error| error.replay) {
5604            self.sql_idempotency_capabilities().await?;
5605            return self
5606                .sql_write_idempotent_once(
5607                    sql,
5608                    &idempotency_key,
5609                    &options,
5610                    fresh_query_id(query_id)?,
5611                    Some(query_id),
5612                )
5613                .await
5614                .map_err(|error| error.error);
5615        }
5616        result.map_err(|error| error.error)
5617    }
5618
5619    async fn sql_write_idempotent_once(
5620        &self,
5621        sql: &str,
5622        idempotency_key: &str,
5623        options: &SqlClientOptions,
5624        query_id: mongreldb_query::QueryId,
5625        expected_original_query_id: Option<mongreldb_query::QueryId>,
5626    ) -> Result<RemoteSqlReceipt, IdempotentAttemptError> {
5627        let timeout_ms = options
5628            .timeout
5629            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
5630        let response = self
5631            .client
5632            .post(self.url("/sql"))
5633            .json(&SqlReq {
5634                sql: sql.to_owned(),
5635                format: None,
5636                query_id,
5637                timeout_ms,
5638                max_output_rows: None,
5639                max_output_bytes: None,
5640                idempotency_key: Some(idempotency_key.to_owned()),
5641                pagination: None,
5642            })
5643            .send()
5644            .await;
5645        let response = match response {
5646            Ok(response) => response,
5647            Err(error) => {
5648                return Err(self
5649                    .idempotent_attempt_loss(query_id, error.to_string())
5650                    .await);
5651            }
5652        };
5653        let response = self
5654            .check_sql_response(response, query_id)
5655            .await
5656            .map_err(IdempotentAttemptError::final_error)?;
5657        let header_error = validate_sql_query_id_header(response.headers(), query_id).err();
5658        let bytes = match bounded_async_bytes(response, MAX_CONTROL_RESPONSE_BYTES).await {
5659            Ok(bytes) => bytes,
5660            Err(error) => return Err(self.idempotent_attempt_loss(query_id, error).await),
5661        };
5662        match decode_remote_sql_receipt(&bytes, query_id, expected_original_query_id) {
5663            Ok(receipt) if header_error.is_none() => Ok(receipt),
5664            Ok(_) => {
5665                let proof = strict_json_value(&bytes).ok().and_then(|value| {
5666                    sql_receipt_commit_proof(&value, query_id, expected_original_query_id)
5667                });
5668                match proof {
5669                    Some(proof) => Err(IdempotentAttemptError::final_error(
5670                        committed_sql_receipt_decode_error(
5671                            query_id,
5672                            proof,
5673                            header_error
5674                                .clone()
5675                                .unwrap_or_else(|| "SQL response identity is invalid".into()),
5676                        ),
5677                    )),
5678                    None => Err(self
5679                        .idempotent_attempt_loss(
5680                            query_id,
5681                            header_error
5682                                .unwrap_or_else(|| "SQL response identity is invalid".into()),
5683                        )
5684                        .await),
5685                }
5686            }
5687            Err(SqlReceiptDecodeError::KnownCommit(error)) => {
5688                Err(IdempotentAttemptError::final_error(error))
5689            }
5690            Err(SqlReceiptDecodeError::Unknown(error)) => {
5691                Err(self.idempotent_attempt_loss(query_id, error).await)
5692            }
5693        }
5694    }
5695
5696    async fn idempotent_attempt_loss(
5697        &self,
5698        query_id: mongreldb_query::QueryId,
5699        message: String,
5700    ) -> IdempotentAttemptError {
5701        let missing = match self
5702            .client
5703            .get(self.url(&format!("/queries/{query_id}")))
5704            .timeout(SQL_RECOVERY_REQUEST_TIMEOUT)
5705            .send()
5706            .await
5707        {
5708            Ok(response) if response.status() == reqwest::StatusCode::NOT_FOUND => {
5709                bounded_async_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
5710                    .await
5711                    .ok()
5712                    .is_some_and(|body| is_exact_query_not_found_response(&body, query_id))
5713            }
5714            _ => false,
5715        };
5716        if missing {
5717            return IdempotentAttemptError {
5718                error: ClientError::QueryOutcomeUnknown {
5719                    query_id: query_id.to_string(),
5720                    message,
5721                    status: None,
5722                    cancel_outcome: None,
5723                },
5724                replay: true,
5725            };
5726        }
5727        IdempotentAttemptError::final_error(
5728            self.recover_after_transport_loss(query_id, message).await,
5729        )
5730    }
5731
5732    pub async fn sql_page(
5733        &self,
5734        sql: &str,
5735        mut options: SqlPageOptions,
5736    ) -> ClientResult<RemoteSqlPage> {
5737        validate_sql_page_options(&options)?;
5738        self.sql_pagination_capabilities().await?;
5739        let query_id = match options.query_id.take() {
5740            Some(query_id) => query_id,
5741            None => mongreldb_query::QueryId::random()
5742                .map_err(|error| ClientError::Transport(error.to_string()))?,
5743        };
5744        let timeout_ms = options
5745            .timeout
5746            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
5747        let response = self
5748            .client
5749            .post(self.url("/sql"))
5750            .json(&SqlReq {
5751                sql: sql.to_owned(),
5752                format: None,
5753                query_id,
5754                timeout_ms,
5755                max_output_rows: options.max_output_rows,
5756                max_output_bytes: options.max_output_bytes,
5757                idempotency_key: None,
5758                pagination: Some(SqlPaginationReq {
5759                    page_size_rows: options.page_size_rows,
5760                    projection: options.projection.clone(),
5761                    max_page_bytes: options.max_page_bytes,
5762                    max_page_tokens: options.max_page_tokens,
5763                }),
5764            })
5765            .send()
5766            .await;
5767        let response = match response {
5768            Ok(response) => response,
5769            Err(error) => {
5770                return Err(self
5771                    .recover_after_transport_loss(query_id, error.to_string())
5772                    .await);
5773            }
5774        };
5775        let response = self.check_sql_response(response, query_id).await?;
5776        if let Err(error) = validate_sql_query_id_header(response.headers(), query_id) {
5777            return Err(self.recover_after_transport_loss(query_id, error).await);
5778        }
5779        let page = bounded_async_bytes(response, MAX_SQL_RESPONSE_BYTES)
5780            .await
5781            .and_then(|bytes| {
5782                strict_json::<RemoteSqlPage>(&bytes, "SQL page").map_err(|error| error.to_string())
5783            })
5784            .and_then(|page| validate_remote_sql_page(page, Some(&options)));
5785        match page {
5786            Ok(page) => Ok(page),
5787            Err(error) => Err(self.recover_after_transport_loss(query_id, error).await),
5788        }
5789    }
5790
5791    pub async fn continue_sql_page(
5792        &self,
5793        cursor: &str,
5794        mut options: RemoteSqlControlOptions,
5795    ) -> ClientResult<RemoteSqlPage> {
5796        if cursor.is_empty() || cursor.len() > 2_048 {
5797            return Err(ClientError::Decode(
5798                "SQL continuation cursor must contain 1 to 2048 bytes".into(),
5799            ));
5800        }
5801        self.sql_pagination_capabilities().await?;
5802        let query_id = match options.query_id.take() {
5803            Some(query_id) => query_id,
5804            None => mongreldb_query::QueryId::random()
5805                .map_err(|error| ClientError::Transport(error.to_string()))?,
5806        };
5807        if options.timeout == Some(std::time::Duration::ZERO) {
5808            return Err(ClientError::Decode("timeout must be positive".into()));
5809        }
5810        let timeout_ms = options
5811            .timeout
5812            .map(|timeout| timeout.as_millis().min(u128::from(u64::MAX)) as u64);
5813        let response = self
5814            .client
5815            .post(self.url("/sql/continue"))
5816            .json(&serde_json::json!({
5817                "cursor": cursor,
5818                "operation_id": query_id,
5819                "timeout_ms": timeout_ms,
5820            }))
5821            .send()
5822            .await
5823            .map_err(|error| ClientError::Transport(error.to_string()))?;
5824        let response = self.check_sql_response(response, query_id).await?;
5825        validate_sql_query_id_header(response.headers(), query_id)
5826            .map_err(|error| client_serialization_error(Some(query_id), error))?;
5827        bounded_async_bytes(response, MAX_SQL_RESPONSE_BYTES)
5828            .await
5829            .and_then(|bytes| {
5830                strict_json::<RemoteSqlPage>(&bytes, "SQL page").map_err(|error| error.to_string())
5831            })
5832            .and_then(|page| validate_remote_sql_page(page, None))
5833            .map_err(|error| client_serialization_error(Some(query_id), error))
5834    }
5835
5836    pub async fn cancel_sql(
5837        &self,
5838        query_id: mongreldb_query::QueryId,
5839    ) -> ClientResult<RemoteCancelOutcome> {
5840        self.sql_cancellation_capabilities().await?;
5841        let response = self
5842            .client
5843            .post(self.url(&format!("/queries/{query_id}/cancel")))
5844            .send()
5845            .await?;
5846        let status = response.status();
5847        if !matches!(
5848            status,
5849            reqwest::StatusCode::OK
5850                | reqwest::StatusCode::ACCEPTED
5851                | reqwest::StatusCode::CONFLICT
5852                | reqwest::StatusCode::NOT_FOUND
5853        ) {
5854            return match self.check(response).await {
5855                Err(error) => Err(error),
5856                Ok(_) => Err(ClientError::Decode(
5857                    "unexpected successful cancellation response".into(),
5858                )),
5859            };
5860        }
5861        let body: serde_json::Value =
5862            decode_async_json(response, MAX_CONTROL_RESPONSE_BYTES, "cancellation").await?;
5863        decode_cancel_outcome(&body, query_id, status.as_u16())
5864    }
5865
5866    pub async fn query_status(
5867        &self,
5868        query_id: mongreldb_query::QueryId,
5869    ) -> ClientResult<RemoteQueryStatus> {
5870        self.sql_cancellation_capabilities().await?;
5871        let response = self
5872            .client
5873            .get(self.url(&format!("/queries/{query_id}")))
5874            .send()
5875            .await?;
5876        let status = decode_async_json(
5877            self.check_query_status_response(response, query_id).await?,
5878            MAX_CONTROL_RESPONSE_BYTES,
5879            "query status",
5880        )
5881        .await?;
5882        validate_remote_query_status(status, query_id).map_err(ClientError::Decode)
5883    }
5884
5885    async fn query_status_optional(
5886        &self,
5887        query_id: mongreldb_query::QueryId,
5888        timeout: std::time::Duration,
5889    ) -> Option<RemoteQueryStatus> {
5890        let response = self
5891            .client
5892            .get(self.url(&format!("/queries/{query_id}")))
5893            .timeout(timeout)
5894            .send()
5895            .await
5896            .ok()?;
5897        if !response.status().is_success() {
5898            return None;
5899        }
5900        let status = decode_async_json(response, MAX_CONTROL_RESPONSE_BYTES, "query status")
5901            .await
5902            .ok()?;
5903        validate_remote_query_status(status, query_id).ok()
5904    }
5905
5906    async fn cancel_sql_for_recovery(
5907        &self,
5908        query_id: mongreldb_query::QueryId,
5909        timeout: std::time::Duration,
5910    ) -> Option<RemoteCancelOutcome> {
5911        let response = self
5912            .client
5913            .post(self.url(&format!("/queries/{query_id}/cancel")))
5914            .timeout(timeout)
5915            .send()
5916            .await
5917            .ok()?;
5918        let status = response.status();
5919        if !matches!(
5920            status,
5921            reqwest::StatusCode::OK
5922                | reqwest::StatusCode::ACCEPTED
5923                | reqwest::StatusCode::CONFLICT
5924                | reqwest::StatusCode::NOT_FOUND
5925        ) {
5926            return None;
5927        }
5928        let body = decode_async_json(response, MAX_CONTROL_RESPONSE_BYTES, "cancellation")
5929            .await
5930            .ok()?;
5931        decode_cancel_outcome(&body, query_id, status.as_u16()).ok()
5932    }
5933
5934    async fn recover_after_transport_loss(
5935        &self,
5936        query_id: mongreldb_query::QueryId,
5937        message: String,
5938    ) -> ClientError {
5939        let deadline = tokio::time::Instant::now() + SQL_RECOVERY_WINDOW;
5940        let mut status = self
5941            .query_status_optional(query_id, SQL_RECOVERY_REQUEST_TIMEOUT)
5942            .await;
5943        if let Some(decisive) = status
5944            .as_ref()
5945            .filter(|status| recovery_status_is_decisive(status))
5946        {
5947            return recovered_query_error(decisive.clone(), message);
5948        }
5949        let cancel_outcome = self
5950            .cancel_sql_for_recovery(
5951                query_id,
5952                deadline
5953                    .saturating_duration_since(tokio::time::Instant::now())
5954                    .min(SQL_RECOVERY_REQUEST_TIMEOUT),
5955            )
5956            .await;
5957        if status.is_none() && cancel_outcome == Some(RemoteCancelOutcome::NotFound) {
5958            return ClientError::QueryOutcomeUnknown {
5959                query_id: query_id.to_string(),
5960                message,
5961                status: None,
5962                cancel_outcome,
5963            };
5964        }
5965        while !status.as_ref().is_some_and(recovery_status_is_decisive)
5966            && tokio::time::Instant::now() < deadline
5967        {
5968            tokio::time::sleep(
5969                deadline
5970                    .saturating_duration_since(tokio::time::Instant::now())
5971                    .min(SQL_RECOVERY_POLL_INTERVAL),
5972            )
5973            .await;
5974            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
5975            if remaining.is_zero() {
5976                break;
5977            }
5978            status = self
5979                .query_status_optional(query_id, remaining.min(SQL_RECOVERY_REQUEST_TIMEOUT))
5980                .await
5981                .or(status);
5982        }
5983        if let Some(decisive) = status
5984            .as_ref()
5985            .filter(|status| recovery_status_is_decisive(status))
5986        {
5987            return recovered_query_error(decisive.clone(), message);
5988        }
5989        ClientError::QueryOutcomeUnknown {
5990            query_id: query_id.to_string(),
5991            message,
5992            status: status.map(Box::new),
5993            cancel_outcome,
5994        }
5995    }
5996
5997    /// Fetch one table's schema + constraint metadata (`GET /kit/schema/{t}`).
5998    pub async fn kit_schema(&self, table: &str) -> ClientResult<TableSchemaInfo> {
5999        let response = self
6000            .client
6001            .get(url_with_segments(
6002                &self.base_url,
6003                &["kit", "schema", table],
6004            )?)
6005            .send()
6006            .await?;
6007        decode_async_json(
6008            self.check(response).await?,
6009            MAX_CONTROL_RESPONSE_BYTES,
6010            "Kit schema",
6011        )
6012        .await
6013    }
6014
6015    /// Run a typed atomic batch (`POST /kit/txn`).
6016    pub async fn kit_txn(&self, req: &KitTxnRequest) -> ClientResult<KitTxnResponse> {
6017        let response = self
6018            .client
6019            .post(self.url("/kit/txn"))
6020            .json(req)
6021            .send()
6022            .await
6023            .map_err(|error| {
6024                kit_txn_outcome_unknown(format!(
6025                    "Kit transaction transport failed before outcome confirmation: {error}"
6026                ))
6027            })?;
6028        if !response.status().is_success() {
6029            let status = response.status().as_u16();
6030            if matches!(status, 401 | 403) {
6031                return Err(kit_txn_auth_error(status));
6032            }
6033            let body = bounded_async_bytes(response, MAX_SQL_RESPONSE_BYTES)
6034                .await
6035                .map_err(|error| {
6036                    kit_txn_outcome_unknown(format!(
6037                        "Kit transaction error response could not be read: {error}"
6038                    ))
6039                })?;
6040            return Err(decode_kit_txn_http_error(status, &body, req));
6041        }
6042        let status = response.status().as_u16();
6043        let body = bounded_async_bytes(response, MAX_SQL_RESPONSE_BYTES)
6044            .await
6045            .map_err(|error| {
6046                kit_txn_outcome_unknown(format!(
6047                    "Kit transaction success response could not be read: {error}"
6048                ))
6049            })?;
6050        decode_kit_txn_success(&body, req, status)
6051    }
6052
6053    /// Run a native typed query (`POST /kit/query`).
6054    pub async fn kit_query(&self, req: &KitQueryRequest) -> ClientResult<KitQueryResponse> {
6055        let response = self
6056            .client
6057            .post(self.url("/kit/query"))
6058            .json(req)
6059            .send()
6060            .await?;
6061        let response = decode_async_json(
6062            self.check(response).await?,
6063            MAX_SQL_RESPONSE_BYTES,
6064            "Kit query",
6065        )
6066        .await?;
6067        validate_kit_query_response(response, req)
6068    }
6069
6070    pub async fn kit_retrieve(
6071        &self,
6072        req: &KitRetrieveRequest,
6073    ) -> ClientResult<KitRetrieveResponse> {
6074        self.kit_retrieve_with_options(req, &AiExecutionOptions::default())
6075            .await
6076    }
6077
6078    pub async fn kit_retrieve_with_options(
6079        &self,
6080        req: &KitRetrieveRequest,
6081        options: &AiExecutionOptions,
6082    ) -> ClientResult<KitRetrieveResponse> {
6083        let response = self
6084            .client
6085            .post(self.url("/kit/retrieve"))
6086            .json(&WithAiExecutionOptions {
6087                request: req,
6088                options,
6089            })
6090            .send()
6091            .await?;
6092        decode_async_json(
6093            self.check(response).await?,
6094            MAX_SQL_RESPONSE_BYTES,
6095            "Kit retrieve",
6096        )
6097        .await
6098    }
6099
6100    pub async fn kit_ann_rerank(
6101        &self,
6102        req: &KitAnnRerankRequest,
6103    ) -> ClientResult<KitAnnRerankResponse> {
6104        self.kit_ann_rerank_with_options(req, &AiExecutionOptions::default())
6105            .await
6106    }
6107
6108    pub async fn kit_ann_rerank_with_options(
6109        &self,
6110        req: &KitAnnRerankRequest,
6111        options: &AiExecutionOptions,
6112    ) -> ClientResult<KitAnnRerankResponse> {
6113        let response = self
6114            .client
6115            .post(self.url("/kit/ann_rerank"))
6116            .json(&WithAiExecutionOptions {
6117                request: req,
6118                options,
6119            })
6120            .send()
6121            .await?;
6122        decode_async_json(
6123            self.check(response).await?,
6124            MAX_SQL_RESPONSE_BYTES,
6125            "Kit ANN rerank",
6126        )
6127        .await
6128    }
6129
6130    pub async fn kit_set_similarity(
6131        &self,
6132        req: &KitSetSimilarityRequest,
6133    ) -> ClientResult<KitSetSimilarityResponse> {
6134        self.kit_set_similarity_with_options(req, &AiExecutionOptions::default())
6135            .await
6136    }
6137
6138    pub async fn kit_set_similarity_with_options(
6139        &self,
6140        req: &KitSetSimilarityRequest,
6141        options: &AiExecutionOptions,
6142    ) -> ClientResult<KitSetSimilarityResponse> {
6143        let response = self
6144            .client
6145            .post(self.url("/kit/set_similarity"))
6146            .json(&WithAiExecutionOptions {
6147                request: req,
6148                options,
6149            })
6150            .send()
6151            .await?;
6152        decode_async_json(
6153            self.check(response).await?,
6154            MAX_SQL_RESPONSE_BYTES,
6155            "Kit set similarity",
6156        )
6157        .await
6158    }
6159
6160    pub async fn kit_search(&self, req: &KitSearchRequest) -> ClientResult<KitSearchResponse> {
6161        let response = self
6162            .client
6163            .post(self.url("/kit/search"))
6164            .json(req)
6165            .send()
6166            .await?;
6167        decode_async_json(
6168            self.check(response).await?,
6169            MAX_SQL_RESPONSE_BYTES,
6170            "Kit search",
6171        )
6172        .await
6173    }
6174
6175    pub async fn kit_ai_metrics(&self) -> ClientResult<serde_json::Value> {
6176        let response = self.client.get(self.url("/kit/ai/metrics")).send().await?;
6177        decode_async_json(
6178            self.check(response).await?,
6179            MAX_CONTROL_RESPONSE_BYTES,
6180            "Kit AI metrics",
6181        )
6182        .await
6183    }
6184}
6185
6186#[derive(Serialize, Clone)]
6187pub struct ColumnDefJson {
6188    pub id: u16,
6189    pub name: String,
6190    pub ty: String,
6191    pub primary_key: bool,
6192}
6193
6194#[derive(Serialize, Clone)]
6195pub struct TxnOp {
6196    pub table: String,
6197    pub op: String,
6198    #[serde(skip_serializing_if = "Option::is_none")]
6199    pub cells: Option<Vec<serde_json::Value>>,
6200    #[serde(skip_serializing_if = "Option::is_none")]
6201    pub row_id: Option<u64>,
6202}
6203
6204fn value_to_json(value: &mongreldb_core::Value) -> ClientResult<serde_json::Value> {
6205    use mongreldb_core::Value;
6206
6207    Ok(match value {
6208        Value::Null => serde_json::Value::Null,
6209        Value::Bool(value) => serde_json::Value::Bool(*value),
6210        Value::Int64(value) => serde_json::Value::Number((*value).into()),
6211        Value::Float64(value) => serde_json::Number::from_f64(*value)
6212            .map(serde_json::Value::Number)
6213            .ok_or_else(|| ClientError::Decode("legacy put rejects non-finite floats".into()))?,
6214        Value::Bytes(value) => tagged_hex_value("bytes", value),
6215        Value::Embedding(values) => serde_json::Value::Array(
6216            values
6217                .iter()
6218                .map(|value| {
6219                    serde_json::Number::from_f64(f64::from(*value))
6220                        .map(serde_json::Value::Number)
6221                        .ok_or_else(|| {
6222                            ClientError::Decode(
6223                                "legacy put rejects non-finite embedding values".into(),
6224                            )
6225                        })
6226                })
6227                .collect::<ClientResult<Vec<_>>>()?,
6228        ),
6229        Value::Decimal(value) => serde_json::json!({
6230            "$mongreldb_type": "decimal",
6231            "unscaled": value.to_string(),
6232        }),
6233        Value::Interval {
6234            months,
6235            days,
6236            nanos,
6237        } => serde_json::json!({
6238            "$mongreldb_type": "interval",
6239            "months": months.to_string(),
6240            "days": days.to_string(),
6241            "nanos": nanos.to_string(),
6242        }),
6243        Value::Uuid(value) => tagged_hex_value("uuid", value),
6244        Value::Json(value) => {
6245            std::str::from_utf8(value)
6246                .map_err(|_| ClientError::Decode("legacy put JSON is not UTF-8".into()))?;
6247            serde_json::from_slice::<serde_json::Value>(value).map_err(|error| {
6248                ClientError::Decode(format!("legacy put JSON is invalid: {error}"))
6249            })?;
6250            tagged_hex_value("json", value)
6251        }
6252    })
6253}
6254
6255fn tagged_hex_value(kind: &str, bytes: &[u8]) -> serde_json::Value {
6256    serde_json::json!({
6257        "$mongreldb_type": kind,
6258        "hex": hex_bytes(bytes),
6259    })
6260}
6261
6262fn hex_bytes(bytes: &[u8]) -> String {
6263    const HEX: &[u8; 16] = b"0123456789abcdef";
6264    let mut encoded = String::with_capacity(bytes.len() * 2);
6265    for byte in bytes {
6266        encoded.push(HEX[(byte >> 4) as usize] as char);
6267        encoded.push(HEX[(byte & 0x0f) as usize] as char);
6268    }
6269    encoded
6270}
6271
6272fn read_arrow_ipc(bytes: &[u8]) -> ClientResult<Vec<RecordBatch>> {
6273    if bytes.is_empty() {
6274        return Ok(Vec::new());
6275    }
6276    let cursor = Cursor::new(bytes);
6277    let reader = FileReader::try_new(cursor, None)
6278        .map_err(|e| ClientError::Decode(format!("arrow ipc: {e}")))?;
6279    reader
6280        .into_iter()
6281        .collect::<std::result::Result<Vec<_>, _>>()
6282        .map_err(|e| ClientError::Decode(format!("arrow read: {e}")))
6283}
6284
6285/// A replication follower that polls a leader's `/wal/stream` endpoint and
6286/// applies WAL records to a local database directory via `SharedWal::replay`.
6287///
6288/// Usage:
6289/// ```no_run
6290/// use mongreldb_client::ReplicationFollower;
6291///
6292/// let mut follower = ReplicationFollower::new("http://leader:8453", "/local/copy")?;
6293/// follower.sync(); // fetches and applies all new records since last sync
6294/// # Ok::<(), mongreldb_client::ClientError>(())
6295/// ```
6296pub struct ReplicationFollower {
6297    leader_url: String,
6298    local_path: std::path::PathBuf,
6299    client: reqwest::blocking::Client,
6300    last_epoch: u64,
6301    bearer_token: Option<SecretString>,
6302    basic_auth: Option<(String, SecretString)>,
6303    local_passphrase: Option<SecretString>,
6304    local_credentials: Option<(String, SecretString)>,
6305}
6306
6307impl std::fmt::Debug for ReplicationFollower {
6308    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6309        formatter
6310            .debug_struct("ReplicationFollower")
6311            .field("leader_url", &self.leader_url)
6312            .field("local_path", &self.local_path)
6313            .field("last_epoch", &self.last_epoch)
6314            .field(
6315                "bearer_token",
6316                &self.bearer_token.as_ref().map(|_| "[REDACTED]"),
6317            )
6318            .field(
6319                "basic_auth",
6320                &self
6321                    .basic_auth
6322                    .as_ref()
6323                    .map(|(username, _)| (username, "[REDACTED]")),
6324            )
6325            .field(
6326                "local_passphrase",
6327                &self.local_passphrase.as_ref().map(|_| "[REDACTED]"),
6328            )
6329            .field(
6330                "local_credentials",
6331                &self
6332                    .local_credentials
6333                    .as_ref()
6334                    .map(|(username, _)| (username, "[REDACTED]")),
6335            )
6336            .finish()
6337    }
6338}
6339
6340impl ReplicationFollower {
6341    /// Create a follower. `leader_url` is the daemon base URL; `local_path` is
6342    /// the local database directory to sync into.
6343    pub fn new(leader_url: &str, local_path: impl AsRef<std::path::Path>) -> ClientResult<Self> {
6344        let leader_url = sanitized_base_url(leader_url).ok_or_else(|| {
6345            ClientError::Transport(
6346                "invalid leader URL: expected http(s) URL without credentials, query, or fragment"
6347                    .into(),
6348            )
6349        })?;
6350        let local_path = local_path.as_ref().to_path_buf();
6351        let last_epoch = mongreldb_core::replica_epoch(&local_path).unwrap_or(0);
6352        Ok(Self {
6353            leader_url,
6354            local_path,
6355            client: reqwest::blocking::Client::new(),
6356            last_epoch,
6357            bearer_token: None,
6358            basic_auth: None,
6359            local_passphrase: None,
6360            local_credentials: None,
6361        })
6362    }
6363
6364    pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
6365        self.bearer_token = Some(token.into().into());
6366        self
6367    }
6368
6369    pub fn with_basic_auth(
6370        mut self,
6371        username: impl Into<String>,
6372        password: impl Into<String>,
6373    ) -> Self {
6374        self.basic_auth = Some((username.into(), password.into().into()));
6375        self
6376    }
6377
6378    pub fn with_local_encryption_passphrase(mut self, passphrase: impl Into<String>) -> Self {
6379        self.local_passphrase = Some(passphrase.into().into());
6380        self
6381    }
6382
6383    pub fn with_local_credentials(
6384        mut self,
6385        username: impl Into<String>,
6386        password: impl Into<String>,
6387    ) -> Self {
6388        self.local_credentials = Some((username.into(), password.into().into()));
6389        self
6390    }
6391
6392    /// Bootstrap when needed, fetch complete committed transactions, append
6393    /// them durably, recover the local database, then advance the watermark.
6394    pub fn sync(&mut self) -> Result<usize, String> {
6395        let legacy_replica_needs_snapshot = mongreldb_core::is_replica(&self.local_path)
6396            && !self.local_path.join("_meta/replication_source_id").exists();
6397        if !self.local_path.join("CATALOG").exists() || legacy_replica_needs_snapshot {
6398            self.ensure_bootstrap_destination_is_owned_or_empty()?;
6399            self.bootstrap()?;
6400        } else if !mongreldb_core::is_replica(&self.local_path) {
6401            return Err(format!(
6402                "refusing to overwrite non-replica database at {}",
6403                self.local_path.display()
6404            ));
6405        }
6406
6407        let mut resp = self.fetch_wal()?;
6408        if resp.status() == reqwest::StatusCode::CONFLICT {
6409            self.validate_snapshot_required_response(resp)?;
6410            self.bootstrap()?;
6411            resp = self.fetch_wal()?;
6412        }
6413        if !resp.status().is_success() {
6414            return Err(format!("leader returned {}", resp.status()));
6415        }
6416        let from_epoch = replication_u64_header(&resp, "x-mongreldb-from-epoch")?;
6417        let leader_epoch = replication_u64_header(&resp, "x-mongreldb-current-epoch")?;
6418        let source_id = replication_digest_header(&resp, "x-mongreldb-source-id")?;
6419        let commit_count = replication_u64_header(&resp, "x-mongreldb-commit-count")?;
6420        let records_sha256 = replication_digest_header(&resp, "x-mongreldb-records-sha256")?;
6421        let earliest_epoch = resp
6422            .headers()
6423            .get("x-mongreldb-earliest-epoch")
6424            .map(|value| {
6425                value
6426                    .to_str()
6427                    .map_err(|error| format!("invalid x-mongreldb-earliest-epoch: {error}"))?
6428                    .parse::<u64>()
6429                    .map_err(|error| format!("invalid x-mongreldb-earliest-epoch: {error}"))
6430            })
6431            .transpose()?;
6432        let body = bounded_blocking_bytes(resp, MAX_REPLICATION_WAL_RESPONSE_BYTES)
6433            .map_err(|error| format!("failed to read WAL response: {error}"))?;
6434        let body = std::str::from_utf8(&body)
6435            .map_err(|_| "invalid WAL response: body is not UTF-8".to_owned())?;
6436        let mut records = Vec::new();
6437        for line in body.lines() {
6438            if line.trim().is_empty() {
6439                continue;
6440            }
6441            records.push(
6442                strict_roundtrip_json::<mongreldb_core::wal::Record>(
6443                    line.as_bytes(),
6444                    "WAL record from leader",
6445                )
6446                .map_err(|error| error.to_string())?,
6447            );
6448        }
6449        let record_count = records.len();
6450        let batch = mongreldb_core::ReplicationBatch::from_wire(
6451            source_id,
6452            from_epoch,
6453            leader_epoch,
6454            earliest_epoch,
6455            commit_count,
6456            records_sha256,
6457            records,
6458        );
6459
6460        let local = self.open_local()?;
6461        let applied_epoch = local
6462            .append_replication_batch(&batch)
6463            .map_err(|error| error.to_string())?;
6464        if record_count == 0 {
6465            return Ok(0);
6466        }
6467        drop(local);
6468        let recovered = self.open_local()?;
6469        if recovered.visible_epoch().0 < applied_epoch {
6470            return Err(format!(
6471                "replica recovery stopped at epoch {}, expected {applied_epoch}",
6472                recovered.visible_epoch().0
6473            ));
6474        }
6475        drop(recovered);
6476        mongreldb_core::write_replica_epoch(&self.local_path, applied_epoch)
6477            .map_err(|error| error.to_string())?;
6478        self.last_epoch = applied_epoch;
6479        Ok(record_count)
6480    }
6481
6482    pub fn bootstrap(&mut self) -> Result<(), String> {
6483        self.ensure_bootstrap_destination_is_owned_or_empty()?;
6484        let url = format!("{}/replication/snapshot", self.leader_url);
6485        let response = self
6486            .request(&url)
6487            .send()
6488            .map_err(|error| format!("failed to fetch replication snapshot: {error}"))?;
6489        if !response.status().is_success() {
6490            return Err(format!(
6491                "leader snapshot endpoint returned {}",
6492                response.status()
6493            ));
6494        }
6495        let response_source_id = replication_digest_header(&response, "x-mongreldb-source-id")?;
6496        let response_epoch = replication_u64_header(&response, "x-mongreldb-current-epoch")?;
6497        let bytes = bounded_blocking_bytes(response, MAX_REPLICATION_SNAPSHOT_BYTES)
6498            .map_err(|error| format!("failed to read replication snapshot: {error}"))?;
6499        let snapshot = mongreldb_core::ReplicationSnapshot::decode(&bytes)
6500            .map_err(|error| error.to_string())?;
6501        if snapshot.source_id() != response_source_id {
6502            return Err("replication snapshot source header does not match snapshot".into());
6503        }
6504        if snapshot.epoch() != response_epoch {
6505            return Err("replication snapshot epoch header does not match snapshot".into());
6506        }
6507        let mut minimum_epoch = self
6508            .last_epoch
6509            .max(mongreldb_core::replica_epoch(&self.local_path).unwrap_or(0));
6510        if self.local_path.join("CATALOG").exists() {
6511            let local = self.open_local()?;
6512            minimum_epoch = minimum_epoch.max(local.visible_epoch().0);
6513        }
6514        if snapshot.epoch() < minimum_epoch {
6515            return Err(format!(
6516                "refusing replication snapshot epoch {} older than local epoch {minimum_epoch}",
6517                snapshot.epoch()
6518            ));
6519        }
6520        snapshot
6521            .install_validated(&self.local_path, |stage| {
6522                let database = self
6523                    .open_path(stage)
6524                    .map_err(mongreldb_core::MongrelError::Other)?;
6525                drop(database);
6526                Ok(())
6527            })
6528            .map_err(|error| error.to_string())?;
6529        self.last_epoch = snapshot.epoch();
6530        Ok(())
6531    }
6532
6533    fn ensure_bootstrap_destination_is_owned_or_empty(&self) -> Result<(), String> {
6534        let metadata = match std::fs::symlink_metadata(&self.local_path) {
6535            Ok(metadata) => metadata,
6536            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()),
6537            Err(error) => {
6538                return Err(format!(
6539                    "failed to inspect replication destination {}: {error}",
6540                    self.local_path.display()
6541                ))
6542            }
6543        };
6544        if metadata.file_type().is_symlink() || !metadata.is_dir() {
6545            return Err(format!(
6546                "replication destination is not a directory: {}",
6547                self.local_path.display()
6548            ));
6549        }
6550        if mongreldb_core::is_replica(&self.local_path) {
6551            return Ok(());
6552        }
6553        let mut entries = std::fs::read_dir(&self.local_path).map_err(|error| {
6554            format!(
6555                "failed to inspect replication destination {}: {error}",
6556                self.local_path.display()
6557            )
6558        })?;
6559        if entries
6560            .next()
6561            .transpose()
6562            .map_err(|error| error.to_string())?
6563            .is_some()
6564        {
6565            return Err(format!(
6566                "refusing to overwrite non-replica directory at {}",
6567                self.local_path.display()
6568            ));
6569        }
6570        drop(entries);
6571        std::fs::remove_dir(&self.local_path).map_err(|error| {
6572            format!(
6573                "failed to claim empty replication destination {}: {error}",
6574                self.local_path.display()
6575            )
6576        })?;
6577        Ok(())
6578    }
6579
6580    fn validate_snapshot_required_response(
6581        &self,
6582        response: reqwest::blocking::Response,
6583    ) -> Result<(), String> {
6584        let from_epoch = replication_u64_header(&response, "x-mongreldb-from-epoch")?;
6585        let current_epoch = replication_u64_header(&response, "x-mongreldb-current-epoch")?;
6586        let source_id = replication_digest_header(&response, "x-mongreldb-source-id")?;
6587        let status = response
6588            .headers()
6589            .get("x-mongreldb-replication-status")
6590            .and_then(|value| value.to_str().ok());
6591        if status != Some("snapshot-required") {
6592            return Err("leader returned an unrecognized replication conflict".into());
6593        }
6594        if self.local_path.join("CATALOG").exists() {
6595            let local_source = mongreldb_core::replica_source_id(&self.local_path)
6596                .map_err(|error| error.to_string())?;
6597            if source_id != local_source {
6598                return Err("replication conflict came from a different source".into());
6599            }
6600        }
6601        replication_u64_header(&response, "x-mongreldb-commit-count")?;
6602        replication_digest_header(&response, "x-mongreldb-records-sha256")?;
6603        if from_epoch != self.last_epoch || current_epoch < self.last_epoch {
6604            return Err(format!(
6605                "invalid replication snapshot requirement: from epoch {from_epoch}, current epoch {current_epoch}, local epoch {}",
6606                self.last_epoch
6607            ));
6608        }
6609        let body = bounded_blocking_bytes(response, MAX_CONTROL_RESPONSE_BYTES)
6610            .map_err(|error| format!("failed to read replication conflict response: {error}"))?;
6611        if body != b"replication snapshot required: WAL retention gap or spilled run" {
6612            return Err("leader returned an unrecognized replication conflict".into());
6613        }
6614        Ok(())
6615    }
6616
6617    fn fetch_wal(&self) -> Result<reqwest::blocking::Response, String> {
6618        let url = format!("{}/wal/stream?since={}", self.leader_url, self.last_epoch);
6619        self.request(&url)
6620            .send()
6621            .map_err(|error| format!("failed to connect to leader: {error}"))
6622    }
6623
6624    fn request(&self, url: &str) -> reqwest::blocking::RequestBuilder {
6625        let request = self.client.get(url);
6626        if let Some(token) = &self.bearer_token {
6627            request.bearer_auth(token.expose_secret())
6628        } else if let Some((username, password)) = &self.basic_auth {
6629            request.basic_auth(username, Some(password.expose_secret()))
6630        } else {
6631            request
6632        }
6633    }
6634
6635    fn open_local(&self) -> Result<mongreldb_core::Database, String> {
6636        self.open_path(&self.local_path)
6637    }
6638
6639    fn open_path(&self, path: &std::path::Path) -> Result<mongreldb_core::Database, String> {
6640        let result = match (&self.local_passphrase, &self.local_credentials) {
6641            (Some(passphrase), Some((username, password))) => {
6642                mongreldb_core::Database::open_encrypted_with_credentials(
6643                    path,
6644                    passphrase.expose_secret(),
6645                    username,
6646                    password.expose_secret(),
6647                )
6648            }
6649            (Some(passphrase), None) => {
6650                mongreldb_core::Database::open_encrypted(path, passphrase.expose_secret())
6651            }
6652            (None, Some((username, password))) => mongreldb_core::Database::open_with_credentials(
6653                path,
6654                username,
6655                password.expose_secret(),
6656            ),
6657            (None, None) => mongreldb_core::Database::open(path),
6658        };
6659        result.map_err(|error| format!("failed to open local replica: {error}"))
6660    }
6661
6662    /// The highest leader commit epoch applied so far.
6663    pub fn last_epoch(&self) -> u64 {
6664        self.last_epoch
6665    }
6666
6667    /// Backward-compatible alias for [`Self::last_epoch`].
6668    pub fn last_seq(&self) -> u64 {
6669        self.last_epoch
6670    }
6671}
6672
6673fn replication_u64_header(
6674    response: &reqwest::blocking::Response,
6675    name: &str,
6676) -> Result<u64, String> {
6677    response
6678        .headers()
6679        .get(name)
6680        .ok_or_else(|| format!("leader response missing {name}"))?
6681        .to_str()
6682        .map_err(|error| format!("invalid {name}: {error}"))?
6683        .parse()
6684        .map_err(|error| format!("invalid {name}: {error}"))
6685}
6686
6687fn replication_digest_header(
6688    response: &reqwest::blocking::Response,
6689    name: &str,
6690) -> Result<[u8; 32], String> {
6691    let value = response
6692        .headers()
6693        .get(name)
6694        .ok_or_else(|| format!("leader response missing {name}"))?
6695        .to_str()
6696        .map_err(|error| format!("invalid {name}: {error}"))?;
6697    if value.len() != 64 {
6698        return Err(format!(
6699            "invalid {name}: expected 64 hexadecimal characters"
6700        ));
6701    }
6702    let mut digest = [0_u8; 32];
6703    for (index, byte) in digest.iter_mut().enumerate() {
6704        *byte = u8::from_str_radix(&value[index * 2..index * 2 + 2], 16)
6705            .map_err(|error| format!("invalid {name}: {error}"))?;
6706    }
6707    Ok(digest)
6708}