1pub 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#[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#[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;
181const 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
420pub 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
432pub 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#[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#[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#[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#[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#[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#[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 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#[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#[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#[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 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 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 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 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"), 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 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 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 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 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 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 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 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 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
6337pub 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 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 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 pub fn last_epoch(&self) -> u64 {
6716 self.last_epoch
6717 }
6718
6719 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}