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