1use std::time::Duration;
16
17use fastrace_reqwest::traceparent_headers;
18use reqwest::IntoUrl;
19use reqwest::Method;
20use reqwest::StatusCode;
21use reqwest::Url;
22use reqwest::header::AUTHORIZATION;
23use reqwest::header::CONTENT_ENCODING;
24use reqwest::header::CONTENT_TYPE;
25use reqwest::header::HeaderValue;
26use serde::Serialize;
27use uuid::Uuid;
28
29use crate::Error;
30use crate::ErrorKind;
31use crate::ResultSet;
32use crate::Statement;
33use crate::Table;
34use crate::protocol::AppendErrorDetails;
35use crate::protocol::AppendRowsResult;
36use crate::protocol::AppendState;
37use crate::protocol::StatementRequest;
38use crate::protocol::StatementStatus;
39use crate::protocol::apply_response_metadata;
40use crate::protocol::decode_statement_response;
41use crate::statement::StatementHandle;
42
43const JSON_CONTENT_ENCODING: &str = "zstd";
44const UNCOMPRESSED_CONTENT_LENGTH: &str = "X-ScopeDB-Uncompressed-Content-Length";
45const ZSTD_COMPRESSION_LEVEL: i32 = 3;
46pub(crate) const MAX_APPEND_BODY_BYTES: usize = 8 * 1024 * 1024;
47pub(crate) const MAX_APPEND_ROWS: usize = 200_000;
48const DEFAULT_WRITE_TIMEOUT: Duration = Duration::from_secs(30);
49
50#[derive(Debug, Clone)]
55pub struct Client {
56 endpoint: Url,
57 client: reqwest::Client,
58 authorization: Option<HeaderValue>,
59 write_timeout: Duration,
60}
61
62pub struct ClientBuilder {
68 endpoint: String,
69 client: Option<reqwest::Client>,
70 api_key: Option<String>,
71 write_timeout: Duration,
72}
73
74impl std::fmt::Debug for ClientBuilder {
75 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 f.debug_struct("ClientBuilder")
77 .field("endpoint", &self.endpoint)
78 .field("client", &self.client)
79 .field("has_api_key", &self.api_key.is_some())
80 .field("write_timeout", &self.write_timeout)
81 .finish()
82 }
83}
84
85impl ClientBuilder {
86 pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
91 self.api_key = Some(api_key.into());
92 self
93 }
94
95 pub fn http_client(mut self, client: reqwest::Client) -> Self {
111 self.client = Some(client);
112 self
113 }
114
115 pub fn write_timeout(mut self, timeout: Duration) -> Self {
129 self.write_timeout = timeout;
130 self
131 }
132
133 pub fn build(self) -> Result<Client, Error> {
135 if self.write_timeout.is_zero() {
136 return Err(Error::new(
137 ErrorKind::ConfigInvalid,
138 "write_timeout must be greater than zero",
139 ));
140 }
141 if tokio::time::Instant::now()
142 .checked_add(self.write_timeout)
143 .is_none()
144 {
145 return Err(Error::new(
146 ErrorKind::ConfigInvalid,
147 "write_timeout is too large",
148 ));
149 }
150 let http_client = match self.client {
151 Some(client) => client,
152 None => reqwest::Client::builder()
153 .redirect(reqwest::redirect::Policy::none())
154 .retry(reqwest::retry::never())
155 .build()
156 .map_err(|err| {
157 Error::new(ErrorKind::ConfigInvalid, "failed to build HTTP client")
158 .set_source(err)
159 })?,
160 };
161 let mut client = Client::new(self.endpoint, http_client)?;
162 client.write_timeout = self.write_timeout;
163 if let Some(api_key) = self.api_key {
164 if api_key.is_empty() {
165 return Err(Error::new(
166 ErrorKind::ConfigInvalid,
167 "api_key must not be empty",
168 ));
169 }
170 let mut authorization =
171 HeaderValue::from_str(&format!("Bearer {api_key}")).map_err(|err| {
172 Error::new(
173 ErrorKind::ConfigInvalid,
174 "api_key is not a valid HTTP credential",
175 )
176 .set_source(err)
177 })?;
178 authorization.set_sensitive(true);
179 client.authorization = Some(authorization);
180 }
181 Ok(client)
182 }
183}
184
185impl Client {
186 pub fn builder(endpoint: impl ToString) -> ClientBuilder {
188 ClientBuilder {
189 endpoint: endpoint.to_string(),
190 client: None,
191 api_key: None,
192 write_timeout: DEFAULT_WRITE_TIMEOUT,
193 }
194 }
195
196 pub fn new<E: IntoUrl>(endpoint: E, client: reqwest::Client) -> Result<Self, Error> {
203 match endpoint.into_url() {
204 Ok(mut endpoint) => {
205 if !matches!(endpoint.scheme(), "http" | "https")
206 || endpoint.host_str().is_none()
207 || !endpoint.username().is_empty()
208 || endpoint.password().is_some()
209 || endpoint.query().is_some()
210 || endpoint.fragment().is_some()
211 {
212 return Err(Error::new(
213 ErrorKind::ConfigInvalid,
214 "endpoint must be an HTTP(S) URL without credentials, query, or fragment",
215 ));
216 }
217 if !endpoint.path().ends_with('/') {
218 let path = format!("{}/", endpoint.path());
219 endpoint.set_path(&path);
220 }
221 Ok(Self {
222 endpoint,
223 client,
224 authorization: None,
225 write_timeout: DEFAULT_WRITE_TIMEOUT,
226 })
227 }
228 Err(err) => Err(Error::new(
229 ErrorKind::ConfigInvalid,
230 "failed to parse endpoint".to_string(),
231 )
232 .set_source(err)),
233 }
234 }
235
236 pub fn statement(&self, statement: impl Into<String>) -> Statement {
237 Statement::new(self.clone(), statement.into())
238 }
239
240 pub fn statement_handle(&self, statement_id: Uuid) -> StatementHandle {
241 StatementHandle::new(self.clone(), statement_id)
242 }
243
244 pub fn table(&self, table: impl Into<String>) -> Table {
245 Table::new(self.clone(), table.into())
246 }
247
248 pub async fn query(&self, statement: impl Into<String>) -> Result<ResultSet, Error> {
252 self.statement(statement).execute().await
253 }
254
255 pub fn catalog(&self) -> crate::Catalog {
257 crate::Catalog::new(self.clone())
258 }
259
260 pub async fn append_rows(
265 &self,
266 database: &str,
267 schema: &str,
268 table: &str,
269 ndjson: impl Into<String>,
270 ) -> Result<AppendRowsResult, Error> {
271 self.append_rows_with_timeout(database, schema, table, ndjson.into(), None)
272 .await
273 }
274
275 pub(crate) async fn append_rows_with_timeout(
276 &self,
277 database: &str,
278 schema: &str,
279 table: &str,
280 ndjson: String,
281 timeout: Option<Duration>,
282 ) -> Result<AppendRowsResult, Error> {
283 if ndjson.len() > MAX_APPEND_BODY_BYTES {
284 return Err(append_rejected_error(format!(
285 "append payload requires {} bytes, exceeds the {MAX_APPEND_BODY_BYTES}-byte append limit",
286 ndjson.len()
287 )));
288 }
289 let expected_rows = ndjson
290 .lines()
291 .filter(|line| !line.trim().is_empty())
292 .count();
293 if expected_rows > MAX_APPEND_ROWS {
294 return Err(append_rejected_error(format!(
295 "append payload exceeds the {MAX_APPEND_ROWS}-row limit",
296 )));
297 }
298 let url = self.make_resource_url(&[
299 "catalog",
300 "databases",
301 database,
302 "schemas",
303 schema,
304 "tables",
305 table,
306 "ingest",
307 ])?;
308 let body =
309 zstd::stream::encode_all(ndjson.as_bytes(), ZSTD_COMPRESSION_LEVEL).map_err(|err| {
310 append_rejected_error("failed to compress table append request body")
311 .set_source(err)
312 })?;
313 let request = self
314 .request(Method::POST, url)
315 .headers(traceparent_headers())
316 .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson")
317 .header(reqwest::header::CONTENT_ENCODING, JSON_CONTENT_ENCODING)
318 .body(body);
319 let mut status = None;
320 let mut headers = reqwest::header::HeaderMap::new();
321 tokio::time::timeout(timeout.unwrap_or(self.write_timeout), async {
322 let response = request.send().await.map_err(|err| {
323 let error = if err.is_timeout() {
324 append_unknown_error("table append request timed out; commit outcome is unknown")
325 .set_write_timeout()
326 } else {
327 append_unknown_error("failed to send table append request")
328 };
329 error.set_source(err)
330 })?;
331 let response_status = response.status();
332 status = Some(response_status);
333 headers = response.headers().clone();
334 let payload = response.bytes().await.map_err(|err| {
335 let error = if err.is_timeout() {
336 append_unknown_error("table append request timed out; commit outcome is unknown")
337 .set_write_timeout()
338 } else {
339 append_unknown_error("failed to read table append response")
340 };
341 error.set_source(err)
342 })?;
343 let result = decode_append_response(response_status, &headers, &payload)?;
344 if result.num_rows_inserted != expected_rows as i64 {
345 return Err(append_unknown_error(format!(
346 "table append response reported {} inserted rows for a {expected_rows}-row request",
347 result.num_rows_inserted
348 )));
349 }
350 Ok(result)
351 })
352 .await
353 .unwrap_or_else(|err| {
354 Err(
355 append_unknown_error("table append request timed out; commit outcome is unknown")
356 .set_write_timeout()
357 .set_source(err),
358 )
359 })
360 .map_err(|err| match status {
361 Some(status) if err.http_status().is_none() => {
362 apply_response_metadata(err, status, &headers)
363 }
364 _ => err,
365 })
366 }
367}
368
369impl Client {
370 #[fastrace::trace]
371 pub(crate) async fn submit_statement(
372 &self,
373 request: StatementRequest,
374 ) -> Result<StatementStatus, Error> {
375 let url = self.make_url("v1/statements")?;
376 let response = self
377 .post_json(url, &request)?
378 .headers(traceparent_headers())
379 .send()
380 .await
381 .map_err(|err| {
382 Error::new(ErrorKind::Unexpected, "failed to submit statement")
383 .set_source(err)
384 .set_persistent()
385 })?;
386 decode_statement_response(response, false)
387 .await
388 .map_err(|error| {
389 if error.http_status().is_none_or(|status| {
392 status.is_success()
393 || status.is_server_error()
394 || status == StatusCode::REQUEST_TIMEOUT
395 }) {
396 error.set_persistent()
397 } else {
398 error
399 }
400 })
401 }
402
403 #[fastrace::trace]
404 pub(crate) async fn fetch_statement(
405 &self,
406 statement_id: Uuid,
407 ) -> Result<StatementStatus, Error> {
408 let path = format!("v1/statements/{statement_id}");
409 let url = self.make_url(&path)?;
410 let response = self
411 .request(Method::GET, url)
412 .headers(traceparent_headers())
413 .query(&[("format", "json")])
414 .send()
415 .await
416 .map_err(|err| {
417 Error::new(
418 ErrorKind::Unexpected,
419 format!("failed to fetch statement: {statement_id:?}"),
420 )
421 .set_source(err)
422 .set_temporary()
423 })?;
424 decode_statement_response(response, true).await
425 }
426
427 #[fastrace::trace]
428 pub(crate) async fn cancel_statement(
429 &self,
430 statement_id: Uuid,
431 ) -> Result<StatementStatus, Error> {
432 let path = format!("v1/statements/{statement_id}/cancel");
433 let url = self.make_url(&path)?;
434 let response = self
435 .request(Method::POST, url)
436 .headers(traceparent_headers())
437 .query(&[("format", "json")])
438 .send()
439 .await
440 .map_err(|err| {
441 Error::new(
442 ErrorKind::Unexpected,
443 format!("failed to cancel statement: {statement_id:?}"),
444 )
445 .set_source(err)
446 .set_temporary()
447 })?;
448 decode_statement_response(response, true).await
449 }
450
451 #[track_caller]
452 fn make_url(&self, path: &str) -> Result<Url, Error> {
453 self.endpoint.join(path).map_err(|err| {
454 Error::new(ErrorKind::Unexpected, "failed to construct URL".to_string()).set_source(err)
455 })
456 }
457
458 pub(crate) fn make_resource_url(&self, segments: &[&str]) -> Result<Url, Error> {
459 if segments.iter().any(|s| matches!(*s, "" | "." | "..")) {
460 return Err(Error::new(
461 ErrorKind::ConfigInvalid,
462 "resource names must not be empty or dot segments",
463 ));
464 }
465 let mut url = self.make_url("v1/")?;
466 let mut path = url.path_segments_mut().map_err(|_| {
467 Error::new(
468 ErrorKind::Unexpected,
469 "endpoint cannot be used as a base URL",
470 )
471 })?;
472 path.pop_if_empty();
473 path.extend(segments.iter().copied());
474 drop(path);
475 Ok(url)
476 }
477
478 pub(crate) fn write_timeout(&self) -> Duration {
479 self.write_timeout
480 }
481
482 pub(crate) fn request(&self, method: Method, url: Url) -> reqwest::RequestBuilder {
483 let request = self
484 .client
485 .request(method, url)
486 .header("X-ScopeDB-Client", "scopedb-client");
487 if let Some(authorization) = &self.authorization {
488 request.header(AUTHORIZATION, authorization.clone())
489 } else {
490 request
491 }
492 }
493
494 fn post_json<T: Serialize>(
495 &self,
496 url: Url,
497 value: &T,
498 ) -> Result<reqwest::RequestBuilder, Error> {
499 let uncompressed = serde_json::to_vec(value).map_err(|err| {
500 Error::new(
501 ErrorKind::Unexpected,
502 "failed to serialize JSON request body",
503 )
504 .set_source(err)
505 })?;
506 let uncompressed_bytes = uncompressed.len();
507 let body = zstd::stream::encode_all(uncompressed.as_slice(), ZSTD_COMPRESSION_LEVEL)
508 .map_err(|err| {
509 Error::new(
510 ErrorKind::Unexpected,
511 "failed to compress JSON request body",
512 )
513 .set_source(err)
514 })?;
515
516 Ok(self
517 .request(Method::POST, url)
518 .header(CONTENT_TYPE, "application/json")
519 .header(CONTENT_ENCODING, JSON_CONTENT_ENCODING)
520 .header(UNCOMPRESSED_CONTENT_LENGTH, uncompressed_bytes.to_string())
521 .body(body))
522 }
523}
524
525fn decode_append_response(
526 status: StatusCode,
527 headers: &reqwest::header::HeaderMap,
528 payload: &[u8],
529) -> Result<AppendRowsResult, Error> {
530 if status == StatusCode::OK {
531 return match serde_json::from_slice::<AppendRowsResult>(payload) {
532 Ok(result)
533 if result.status == AppendState::Committed && result.num_rows_inserted >= 0 =>
534 {
535 Ok(result)
536 }
537 Ok(_) => Err(apply_response_metadata(
538 append_unknown_error("table append response has an invalid body"),
539 status,
540 headers,
541 )),
542 Err(err) => Err(apply_response_metadata(
543 append_unknown_error("failed to decode table append response").set_source(err),
544 status,
545 headers,
546 )),
547 };
548 }
549
550 if status.is_success() {
551 return Err(apply_response_metadata(
552 append_unknown_error(format!(
553 "table append returned unexpected success status {}",
554 status.as_u16()
555 )),
556 status,
557 headers,
558 ));
559 }
560
561 let error =
562 crate::protocol::decode_http_error(status, headers, payload, ErrorKind::AppendRowsFailed);
563 if let Some(details) = error
564 .details()
565 .and_then(|value| serde_json::from_value::<AppendErrorDetails>(value.clone()).ok())
566 && matches!(details.status, AppendState::Rejected | AppendState::Unknown)
567 {
568 let status = details.status;
569 let error = error.set_append_details(details);
570
571 return Err(if status == AppendState::Unknown {
572 error.set_persistent()
573 } else {
574 error
575 });
576 }
577
578 Err(error
579 .set_append_details(AppendErrorDetails {
580 status: AppendState::Unknown,
581 row_errors: Vec::new(),
582 row_errors_truncated: false,
583 })
584 .set_persistent())
585}
586
587fn append_unknown_error(message: impl Into<String>) -> Error {
588 Error::new(ErrorKind::AppendRowsFailed, message)
589 .set_append_details(AppendErrorDetails {
590 status: AppendState::Unknown,
591 row_errors: Vec::new(),
592 row_errors_truncated: false,
593 })
594 .set_persistent()
595}
596
597fn append_rejected_error(message: impl Into<String>) -> Error {
598 Error::new(ErrorKind::AppendRowsFailed, message).set_append_details(AppendErrorDetails {
599 status: AppendState::Rejected,
600 row_errors: Vec::new(),
601 row_errors_truncated: false,
602 })
603}
604
605#[cfg(test)]
606mod tests {
607 use reqwest::StatusCode;
608 use reqwest::header::HeaderMap;
609 use serde_json::json;
610
611 use super::Client;
612 use super::JSON_CONTENT_ENCODING;
613 use super::UNCOMPRESSED_CONTENT_LENGTH;
614 use super::decode_append_response;
615 use crate::ErrorKind;
616 use crate::protocol::AppendState;
617
618 fn decode(status: StatusCode, payload: &[u8]) -> Result<crate::AppendRowsResult, crate::Error> {
619 decode_append_response(status, &HeaderMap::new(), payload)
620 }
621
622 #[test]
623 fn resource_url_preserves_base_path_and_encodes_segments() {
624 let client = Client::new("https://example.com/proxy", reqwest::Client::new()).unwrap();
625
626 let url = client
627 .make_resource_url(&["databases", "analytics/2026", "schemas", "events archive"])
628 .unwrap();
629
630 assert_eq!(
631 url.as_str(),
632 "https://example.com/proxy/v1/databases/analytics%2F2026/schemas/events%20archive"
633 );
634 }
635
636 #[test]
637 fn client_builder_validates_and_redacts_api_keys() {
638 let builder = Client::builder("https://example.com").api_key("secret-api-key");
639 assert!(!format!("{builder:?}").contains("secret-api-key"));
640 let client = Client::builder("https://example.com")
641 .api_key("secret-api-key")
642 .build()
643 .unwrap();
644 assert!(!format!("{client:?}").contains("secret-api-key"));
645
646 let error = Client::builder("https://example.com")
647 .api_key("")
648 .build()
649 .unwrap_err();
650 assert_eq!(error.kind(), ErrorKind::ConfigInvalid);
651 }
652
653 #[test]
654 fn write_timeout_defaults_and_validation() {
655 let client = Client::new("https://example.com", reqwest::Client::new()).unwrap();
656 assert_eq!(client.write_timeout, std::time::Duration::from_secs(30));
657 let client = Client::builder("https://example.com").build().unwrap();
658 assert_eq!(client.write_timeout, std::time::Duration::from_secs(30));
659 for timeout in [std::time::Duration::ZERO, std::time::Duration::MAX] {
660 let error = Client::builder("https://example.com")
661 .write_timeout(timeout)
662 .build()
663 .unwrap_err();
664 assert_eq!(error.kind(), ErrorKind::ConfigInvalid);
665 }
666 }
667
668 #[test]
669 fn json_requests_use_zstd_by_default() {
670 let client = Client::new("https://example.com", reqwest::Client::new()).unwrap();
671 let value = json!({"statement": "SELECT 1", "format": "json"});
672 let expected = serde_json::to_vec(&value).unwrap();
673 let request = client
674 .post_json(client.make_url("v1/statements").unwrap(), &value)
675 .unwrap()
676 .build()
677 .unwrap();
678
679 assert_eq!(
680 request
681 .headers()
682 .get(reqwest::header::CONTENT_ENCODING)
683 .unwrap(),
684 JSON_CONTENT_ENCODING
685 );
686 assert_eq!(
687 request
688 .headers()
689 .get(UNCOMPRESSED_CONTENT_LENGTH)
690 .unwrap()
691 .to_str()
692 .unwrap(),
693 expected.len().to_string()
694 );
695 let compressed = request.body().unwrap().as_bytes().unwrap();
696 assert_eq!(zstd::stream::decode_all(compressed).unwrap(), expected);
697 }
698
699 #[test]
700 fn committed_append_response_is_returned() {
701 let result = decode(
702 StatusCode::OK,
703 br#"{"status":"committed","num_rows_inserted":2}"#,
704 )
705 .unwrap();
706
707 assert_eq!(result.status, AppendState::Committed);
708 assert_eq!(result.num_rows_inserted, 2);
709 }
710
711 #[test]
712 fn rejected_append_error_preserves_row_details_and_retry_status() {
713 let error = decode(
714 StatusCode::SERVICE_UNAVAILABLE,
715 br#"{"message":"validation failed","code":"ingest_failure","details":{"status":"rejected","row_errors":[{"row_index":3,"column":"ts","message":"invalid timestamp"}],"row_errors_truncated":true}}"#,
716 )
717 .unwrap_err();
718
719 assert_eq!(error.kind(), ErrorKind::AppendRowsFailed);
720 assert!(error.is_temporary());
721 let details = error.append_details().unwrap();
722 assert_eq!(details.status, AppendState::Rejected);
723 assert_eq!(details.row_errors[0].row_index, 3);
724 assert!(details.row_errors_truncated);
725 }
726
727 #[test]
728 fn append_error_preserves_server_message_and_http_metadata() {
729 let mut headers = HeaderMap::new();
730 headers.insert("x-request-id", "header-request".parse().unwrap());
731 headers.insert(reqwest::header::RETRY_AFTER, "3".parse().unwrap());
732 let error = decode_append_response(
733 StatusCode::SERVICE_UNAVAILABLE,
734 &headers,
735 br#"{"message":"capacity is temporarily unavailable","request_id":"body-request","retryable":false,"code":"ingest_failure","details":{"status":"rejected"}}"#,
736 )
737 .unwrap_err();
738
739 assert_eq!(error.message(), "capacity is temporarily unavailable");
740 assert_eq!(error.http_status(), Some(StatusCode::SERVICE_UNAVAILABLE));
741 assert_eq!(error.request_id(), Some("body-request"));
742 assert_eq!(error.retry_after(), Some(std::time::Duration::from_secs(3)));
743 assert!(!error.is_retryable());
744 }
745
746 #[test]
747 fn unknown_append_outcome_is_never_retryable() {
748 let error = decode(
749 StatusCode::SERVICE_UNAVAILABLE,
750 br#"{"message":"coordinator unavailable","code":"ingest_failure","details":{"status":"unknown"}}"#,
751 )
752 .unwrap_err();
753
754 assert!(error.is_persistent());
755 assert!(!error.is_temporary());
756 assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
757 }
758
759 #[test]
760 fn unstructured_append_error_is_treated_as_unknown() {
761 let error = decode(StatusCode::BAD_REQUEST, br#"{"message":"bad request"}"#).unwrap_err();
762
763 assert!(error.is_persistent());
764 assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
765 }
766
767 #[test]
768 fn nested_append_error_has_clean_message_and_unknown_outcome() {
769 let error = decode(
770 StatusCode::NOT_FOUND,
771 br#"{"error":{"message":"unsupported path"}}"#,
772 )
773 .unwrap_err();
774
775 assert_eq!(error.message(), "unsupported path");
776 assert!(error.is_persistent());
777 assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
778 }
779
780 #[test]
781 fn nested_append_rejection_keeps_its_outcome_and_diagnostics() {
782 let error = decode(
783 StatusCode::SERVICE_UNAVAILABLE,
784 br#"{"request_id":"proxy-request","error":{"message":"try later","code":"ingest_failure","retryable":false,"details":{"status":"rejected","row_errors":[{"row_index":1,"column":"id","message":"invalid int"}]}}}"#,
785 )
786 .unwrap_err();
787 assert_eq!(error.code(), Some("ingest_failure"));
788 assert_eq!(error.request_id(), Some("proxy-request"));
789 assert_eq!(error.http_status(), Some(StatusCode::SERVICE_UNAVAILABLE));
790 assert!(!error.is_retryable());
791 let details = error.append_details().unwrap();
792 assert_eq!(details.status, AppendState::Rejected);
793 assert_eq!(details.row_errors[0].row_index, 1);
794 }
795}