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