1#[cfg(feature = "recording")]
52use std::collections::HashMap;
53#[cfg(feature = "recording")]
54use std::path::PathBuf;
55use std::sync::Arc;
56#[cfg(feature = "recording")]
57use std::sync::atomic::{AtomicU64, Ordering};
58use std::time::Duration;
59
60use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
61use reqwest::{Body, Client, IntoUrl, Method, RequestBuilder};
62use serde::Serialize;
63use tokio::runtime::Handle;
64
65use self::retry::RetryExt;
66use self::service::{HttpService, make_service};
67pub use self::token::{TemporaryToken, TokenCache};
68
69pub mod aws;
70pub mod azure;
71mod backoff;
72mod client;
73mod config;
74#[cfg(feature = "connectrpc")]
75pub mod connectrpc;
76mod credential;
77pub mod databricks;
78mod error;
79pub mod gcp;
80pub mod service;
81
82mod retry;
83mod token;
84mod util;
85
86pub use client::{Certificate, ClientConfigKey, ClientOptions};
87pub use credential::*;
88pub use error::*;
89pub use retry::Error as RetryError;
90pub use retry::RetryConfig;
91pub use service::{ReqwestService, SpawnService};
92
93type SharedSigner = Arc<dyn RequestSigner>;
95
96#[derive(Debug)]
98struct NoopSigner;
99
100impl RequestSigner for NoopSigner {
101 fn sign<'a>(
102 &'a self,
103 req: RequestBuilder,
104 ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
105 Box::pin(async move { Ok(req) })
106 }
107}
108
109#[derive(Debug)]
111struct BearerTokenSigner {
112 token: String,
113}
114
115impl RequestSigner for BearerTokenSigner {
116 fn sign<'a>(
117 &'a self,
118 req: RequestBuilder,
119 ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
120 let token = self.token.clone();
121 Box::pin(async move { Ok(req.bearer_auth(&token)) })
122 }
123}
124
125#[cfg(feature = "recording")]
126#[derive(Debug, Clone)]
127struct RecordingState {
128 out_dir: PathBuf,
129 counter: Arc<AtomicU64>,
130}
131
132#[derive(Clone)]
139pub struct CloudClient {
140 signer: SharedSigner,
141 reqwest_client: Client,
142 service: Arc<dyn HttpService>,
143 pub retry_config: RetryConfig,
150 #[cfg(feature = "recording")]
151 recording: Option<RecordingState>,
152}
153
154impl CloudClient {
155 fn new_with_signer(
156 signer: SharedSigner,
157 reqwest_client: Client,
158 service: Arc<dyn HttpService>,
159 retry_config: RetryConfig,
160 ) -> Self {
161 Self {
162 signer,
163 reqwest_client,
164 service,
165 retry_config,
166 #[cfg(feature = "recording")]
167 recording: None,
168 }
169 }
170
171 pub fn new_aws<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
176 where
177 I: IntoIterator<Item = (K, V)>,
178 K: AsRef<str>,
179 V: Into<String>,
180 {
181 let config = options
182 .into_iter()
183 .fold(
184 aws::AmazonBuilder::new(),
185 |builder, (key, value)| match key.as_ref().parse() {
186 Ok(k) => builder.with_config(k, value),
187 Err(_) => builder,
188 },
189 )
190 .build(runtime)?;
191
192 let reqwest_client = config.client_options.client()?;
193 let service = make_service(reqwest_client.clone(), runtime);
194 let retry_config = config.retry_config.clone();
195 Ok(Self::new_with_signer(
196 Arc::new(config),
197 reqwest_client,
198 service,
199 retry_config,
200 ))
201 }
202
203 pub fn new_google<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
208 where
209 I: IntoIterator<Item = (K, V)>,
210 K: AsRef<str>,
211 V: Into<String>,
212 {
213 let config = options
214 .into_iter()
215 .fold(
216 gcp::GoogleBuilder::new(),
217 |builder, (key, value)| match key.as_ref().parse() {
218 Ok(k) => builder.with_config(k, value),
219 Err(_) => builder,
220 },
221 )
222 .build(runtime)?;
223
224 let reqwest_client = config.client_options.client()?;
225 let service = make_service(reqwest_client.clone(), runtime);
226 let retry_config = config.retry_config.clone();
227 Ok(Self::new_with_signer(
228 Arc::new(config),
229 reqwest_client,
230 service,
231 retry_config,
232 ))
233 }
234
235 pub fn new_azure<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
240 where
241 I: IntoIterator<Item = (K, V)>,
242 K: AsRef<str>,
243 V: Into<String>,
244 {
245 let config = options
246 .into_iter()
247 .fold(
248 azure::AzureBuilder::new(),
249 |builder, (key, value)| match key.as_ref().parse() {
250 Ok(k) => builder.with_config(k, value),
251 Err(_) => builder,
252 },
253 )
254 .build(runtime)?;
255
256 let reqwest_client = config.client_options.client()?;
257 let service = make_service(reqwest_client.clone(), runtime);
258 let retry_config = config.retry_config.clone();
259 Ok(Self::new_with_signer(
260 Arc::new(config),
261 reqwest_client,
262 service,
263 retry_config,
264 ))
265 }
266
267 pub fn new_databricks<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
279 where
280 I: IntoIterator<Item = (K, V)>,
281 K: AsRef<str>,
282 V: Into<String>,
283 {
284 let builder = options.into_iter().fold(
285 databricks::DatabricksBuilder::new(),
286 |builder, (key, value)| match key.as_ref().parse() {
287 Ok(k) => builder.with_config(k, value),
288 Err(_) => builder,
289 },
290 );
291 Self::new_databricks_from_builder(builder, runtime)
292 }
293
294 pub fn new_databricks_from_builder(
305 builder: databricks::DatabricksBuilder,
306 runtime: Option<&Handle>,
307 ) -> Result<Self> {
308 let config = builder.build(runtime)?;
309
310 let reqwest_client = config.client_options.client()?;
311 let service = make_service(reqwest_client.clone(), runtime);
312 let retry_config = config.retry_config.clone();
313 Ok(Self::new_with_signer(
314 Arc::new(config),
315 reqwest_client,
316 service,
317 retry_config,
318 ))
319 }
320
321 pub fn new_with_token(token: impl ToString) -> Self {
323 let reqwest_client = Client::new();
324 let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
325 Self::new_with_signer(
326 Arc::new(BearerTokenSigner {
327 token: token.to_string(),
328 }),
329 reqwest_client,
330 service,
331 RetryConfig::default(),
332 )
333 }
334
335 pub fn new_unauthenticated() -> Self {
337 let reqwest_client = Client::new();
338 let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
339 Self::new_with_signer(
340 Arc::new(NoopSigner),
341 reqwest_client,
342 service,
343 RetryConfig::default(),
344 )
345 }
346
347 pub fn with_runtime(mut self, handle: Handle) -> Self {
353 self.service = Arc::new(SpawnService::new(self.service, handle));
354 self
355 }
356
357 pub fn with_http_service(mut self, service: Arc<dyn HttpService>) -> Self {
359 self.service = service;
360 self
361 }
362
363 pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
371 self.retry_config = retry_config;
372 self
373 }
374
375 pub async fn sign_and_send(&self, request: reqwest::Request) -> Result<reqwest::Response> {
393 let builder = self
397 .reqwest_client
398 .request(request.method().clone(), request.url().clone());
399 let builder = builder.headers(request.headers().clone());
400 let builder = match request.body().and_then(|b| b.as_bytes()) {
401 Some(bytes) => builder.body(bytes.to_vec()),
402 None => builder,
403 };
404 let builder = self.signer.sign(builder).await?;
405 builder
406 .send_retry(&self.retry_config, self.service.clone())
407 .await
408 .map_err(|e| e.error())
409 }
410
411 pub fn request<U: IntoUrl>(&self, method: Method, url: U) -> CloudRequestBuilder {
417 CloudRequestBuilder {
418 builder: self.reqwest_client.request(method, url),
419 client: self.clone(),
420 #[cfg(feature = "recording")]
421 out_dir: self.recording.as_ref().map(|r| r.out_dir.clone()),
422 }
423 }
424
425 pub fn get<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
427 self.request(Method::GET, url)
428 }
429
430 pub fn post<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
432 self.request(Method::POST, url)
433 }
434
435 pub fn put<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
437 self.request(Method::PUT, url)
438 }
439
440 pub fn delete<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
442 self.request(Method::DELETE, url)
443 }
444
445 pub fn head<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
447 self.request(Method::HEAD, url)
448 }
449
450 pub fn patch<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
452 self.request(Method::PATCH, url)
453 }
454
455 pub fn options<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
457 self.request(Method::OPTIONS, url)
458 }
459
460 pub fn trace<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
462 self.request(Method::TRACE, url)
463 }
464
465 pub fn connect<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
467 self.request(Method::CONNECT, url)
468 }
469
470 #[cfg(feature = "recording")]
489 pub fn set_recording_dir(&mut self, out_dir: std::path::PathBuf) -> Result<(), std::io::Error> {
490 let out_dir = std::fs::canonicalize(out_dir)?;
491 self.recording = Some(RecordingState {
492 out_dir,
493 counter: Arc::new(AtomicU64::new(0)),
494 });
495 Ok(())
496 }
497}
498
499pub struct CloudRequestBuilder {
505 builder: RequestBuilder,
506 client: CloudClient,
507 #[cfg(feature = "recording")]
508 out_dir: Option<PathBuf>,
509}
510
511impl CloudRequestBuilder {
512 pub fn header<K, V>(mut self, key: K, value: V) -> CloudRequestBuilder
514 where
515 HeaderName: TryFrom<K>,
516 <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
517 HeaderValue: TryFrom<V>,
518 <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
519 {
520 self.builder = self.builder.header(key, value);
521 self
522 }
523
524 pub fn headers(mut self, headers: HeaderMap) -> CloudRequestBuilder {
528 self.builder = self.builder.headers(headers);
529 self
530 }
531
532 pub fn body<T: Into<Body>>(mut self, body: T) -> CloudRequestBuilder {
534 self.builder = self.builder.body(body);
535 self
536 }
537
538 pub fn timeout(mut self, timeout: Duration) -> CloudRequestBuilder {
544 self.builder = self.builder.timeout(timeout);
545 self
546 }
547
548 pub fn query<T: Serialize + ?Sized>(mut self, query: &T) -> CloudRequestBuilder {
567 self.builder = self.builder.query(query);
568 self
569 }
570
571 pub fn json<T: Serialize + ?Sized>(mut self, json: &T) -> CloudRequestBuilder {
578 self.builder = self.builder.json(json);
579 self
580 }
581
582 pub async fn send(mut self) -> Result<reqwest::Response> {
605 self.builder = self.client.signer.sign(self.builder).await?;
606
607 #[cfg(not(feature = "recording"))]
608 {
609 self.builder
610 .send_retry(&self.client.retry_config, self.client.service.clone())
611 .await
612 .map_err(|e| e.error())
613 }
614 #[cfg(feature = "recording")]
615 {
616 let response = send_record(self).await?;
617 Ok(response)
618 }
619 }
620
621 #[cfg(not(feature = "recording"))]
648 pub async fn send_raw(mut self) -> Result<reqwest::Response, SendRawError> {
649 self.builder = self.client.signer.sign(self.builder).await?;
650 self.builder
651 .send_retry(&self.client.retry_config, self.client.service.clone())
652 .await
653 .map_err(SendRawError::Retry)
654 }
655
656 #[cfg(feature = "recording")]
659 pub async fn send_raw(self) -> Result<reqwest::Response, SendRawError> {
660 send_record(self).await
661 }
662}
663
664#[cfg(not(feature = "recording"))]
672#[derive(Debug, thiserror::Error)]
673pub enum SendRawError {
674 #[error(transparent)]
676 Sign(#[from] Error),
677 #[error(transparent)]
679 Retry(RetryError),
680}
681
682#[cfg(feature = "recording")]
685pub type SendRawError = Error;
686
687#[cfg(feature = "recording")]
688#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
689pub struct RequestResponseInfo {
690 pub request: RequestInfo,
691 pub response: ResponseInfo,
692}
693
694#[cfg(feature = "recording")]
695#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
696pub struct RequestInfo {
697 pub method: String,
698 pub url_path: String,
699 pub body: Option<String>,
700}
701
702#[cfg(feature = "recording")]
703#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
704pub struct ResponseInfo {
705 pub status: u16,
706 pub headers: HashMap<String, String>,
707 pub body: Option<String>,
708}
709
710#[cfg(feature = "recording")]
715const SENSITIVE_HEADERS: &[&str] = &[
716 "authorization",
717 "x-amz-security-token",
718 "x-amz-content-sha256",
719 "x-databricks-authorization",
720 "x-ms-identity-principal-id",
721 "x-goog-iam-credentials-token",
722 "cookie",
723 "set-cookie",
724];
725
726#[cfg(feature = "recording")]
727fn redact_headers(headers: &HashMap<String, String>) -> HashMap<String, String> {
728 headers
729 .iter()
730 .map(|(k, v)| {
731 let v = if SENSITIVE_HEADERS.contains(&k.to_lowercase().as_str()) {
732 "<REDACTED>".to_string()
733 } else {
734 v.clone()
735 };
736 (k.clone(), v)
737 })
738 .collect()
739}
740
741#[cfg(feature = "recording")]
742async fn send_record(builder: CloudRequestBuilder) -> Result<reqwest::Response> {
743 let Some(out_dir) = builder.out_dir else {
744 let request = builder.builder.build().expect("request to be valid");
745 return builder.client.service.call(request).await;
746 };
747 let (_client, request) = builder.builder.build_split();
748 let request = request.expect("request to be valid");
749
750 let request_info = RequestInfo {
751 method: request.method().as_str().to_string(),
752 url_path: {
753 let url = request.url();
754 match url.query() {
755 Some(query) => format!("{}?{}", url.path(), query),
756 None => url.path().to_string(),
757 }
758 },
759 body: request
760 .body()
761 .and_then(|b| b.as_bytes().map(|b| String::from_utf8_lossy(b).to_string())),
762 };
763
764 let response = builder.client.service.call(request).await?;
765
766 let status = response.status().as_u16();
768 let raw_headers: HashMap<String, String> = response
769 .headers()
770 .iter()
771 .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
772 .collect();
773
774 let response_bytes = response.bytes().await?;
776 let response_body = if response_bytes.is_empty() {
777 None
778 } else {
779 Some(String::from_utf8_lossy(&response_bytes).to_string())
780 };
781
782 let recording = RequestResponseInfo {
783 request: request_info,
784 response: ResponseInfo {
785 status,
786 headers: redact_headers(&raw_headers),
788 body: response_body,
789 },
790 };
791
792 let counter = builder
793 .client
794 .recording
795 .as_ref()
796 .map(|r| r.counter.fetch_add(1, Ordering::SeqCst))
797 .unwrap_or(0);
798 let file_path = out_dir.join(format!("{counter:04}.json"));
799 if let Err(e) = std::fs::File::create(&file_path)
800 .and_then(|f| serde_json::to_writer_pretty(f, &recording).map_err(Into::into))
801 {
802 tracing::warn!(
803 "Failed to write recording to {}: {}",
804 file_path.display(),
805 e
806 );
807 }
808
809 let mut mock_response = http::Response::builder().status(status);
811 for (k, v) in &raw_headers {
812 mock_response = mock_response.header(k, v);
813 }
814 let mock_response = mock_response
815 .body(response_bytes)
816 .expect("valid status code and headers");
817
818 Ok(reqwest::Response::from(mock_response))
819}
820
821#[cfg(all(test, feature = "recording"))]
822mod tests {
823 use super::*;
824 use std::fs;
825 use tempfile::TempDir;
826
827 #[tokio::test]
828 async fn test_request_response_recording() {
829 let temp_dir = TempDir::new().unwrap();
831 let temp_path = temp_dir.path().to_path_buf();
832
833 let mut server = mockito::Server::new_async().await;
835 let mock = server
836 .mock("GET", "/test")
837 .with_status(200)
838 .with_header("content-type", "application/json")
839 .with_body(r#"{"message": "Hello, World!"}"#)
840 .create_async()
841 .await;
842
843 let mut client = CloudClient::new_unauthenticated();
845 client.set_recording_dir(temp_path.clone()).unwrap();
846
847 let url = format!("{}/test", server.url());
849 let response = client.get(&url).send().await.unwrap();
850
851 assert_eq!(response.status(), 200);
853 let body = response.text().await.unwrap();
854 assert_eq!(body, r#"{"message": "Hello, World!"}"#);
855
856 let recordings: Vec<_> = fs::read_dir(&temp_path)
858 .unwrap()
859 .filter_map(|entry| {
860 let entry = entry.ok()?;
861 let path = entry.path();
862 if path.extension()? == "json" {
863 Some(path)
864 } else {
865 None
866 }
867 })
868 .collect();
869
870 assert_eq!(recordings.len(), 1, "Expected exactly one recording file");
871
872 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
874 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
875
876 assert_eq!(recording.request.method, "GET");
878 assert_eq!(recording.request.url_path, "/test");
879 assert_eq!(recording.request.body, None);
880
881 assert_eq!(recording.response.status, 200);
883 assert_eq!(
884 recording.response.headers.get("content-type").unwrap(),
885 "application/json"
886 );
887 assert_eq!(
888 recording.response.body.as_ref().unwrap(),
889 r#"{"message": "Hello, World!"}"#
890 );
891
892 mock.assert_async().await;
893 }
894
895 #[tokio::test]
896 async fn test_recording_with_request_body() {
897 let temp_dir = TempDir::new().unwrap();
899 let temp_path = temp_dir.path().to_path_buf();
900
901 let mut server = mockito::Server::new_async().await;
903 let mock = server
904 .mock("POST", "/create")
905 .with_status(201)
906 .with_header("location", "/resource/123")
907 .with_body(r#"{"id": 123, "status": "created"}"#)
908 .create_async()
909 .await;
910
911 let mut client = CloudClient::new_unauthenticated();
913 client.set_recording_dir(temp_path.clone()).unwrap();
914
915 let url = format!("{}/create", server.url());
917 let response = client
918 .post(&url)
919 .json(&serde_json::json!({"name": "test resource"}))
920 .send()
921 .await
922 .unwrap();
923
924 assert_eq!(response.status(), 201);
926
927 let recordings: Vec<_> = fs::read_dir(&temp_path)
929 .unwrap()
930 .filter_map(|entry| {
931 let entry = entry.ok()?;
932 let path = entry.path();
933 if path.extension()? == "json" {
934 Some(path)
935 } else {
936 None
937 }
938 })
939 .collect();
940
941 assert_eq!(recordings.len(), 1);
942
943 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
945 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
946
947 assert_eq!(recording.request.method, "POST");
949 assert_eq!(recording.request.url_path, "/create");
950 assert!(recording.request.body.is_some());
951 assert!(recording.request.body.unwrap().contains("test resource"));
952
953 assert_eq!(recording.response.status, 201);
955 assert_eq!(
956 recording.response.headers.get("location").unwrap(),
957 "/resource/123"
958 );
959 assert!(recording.response.body.unwrap().contains("created"));
960
961 mock.assert_async().await;
962 }
963
964 #[tokio::test]
965 async fn test_counter_based_file_naming() {
966 let temp_dir = TempDir::new().unwrap();
968 let temp_path = temp_dir.path().to_path_buf();
969
970 let mut server = mockito::Server::new_async().await;
972 let mock1 = server
973 .mock("GET", "/first")
974 .with_status(200)
975 .with_body("first response")
976 .create_async()
977 .await;
978 let mock2 = server
979 .mock("GET", "/second")
980 .with_status(200)
981 .with_body("second response")
982 .create_async()
983 .await;
984
985 let mut client = CloudClient::new_unauthenticated();
987 client.set_recording_dir(temp_path.clone()).unwrap();
988
989 let url1 = format!("{}/first", server.url());
991 let url2 = format!("{}/second", server.url());
992
993 let _response1 = client.get(&url1).send().await.unwrap();
994 let _response2 = client.get(&url2).send().await.unwrap();
995
996 let mut recordings: Vec<_> = fs::read_dir(&temp_path)
998 .unwrap()
999 .filter_map(|entry| {
1000 let entry = entry.ok()?;
1001 let path = entry.path();
1002 if path.extension()? == "json" {
1003 Some(path)
1004 } else {
1005 None
1006 }
1007 })
1008 .collect();
1009
1010 recordings.sort();
1011 assert_eq!(recordings.len(), 2);
1012
1013 assert!(recordings[0].file_name().unwrap().to_str().unwrap() == "0000.json");
1015 assert!(recordings[1].file_name().unwrap().to_str().unwrap() == "0001.json");
1016
1017 let first_content = fs::read_to_string(&recordings[0]).unwrap();
1019 let first_recording: RequestResponseInfo = serde_json::from_str(&first_content).unwrap();
1020 assert_eq!(first_recording.request.url_path, "/first");
1021 assert_eq!(
1022 first_recording.response.body.as_ref().unwrap(),
1023 "first response"
1024 );
1025
1026 let second_content = fs::read_to_string(&recordings[1]).unwrap();
1027 let second_recording: RequestResponseInfo = serde_json::from_str(&second_content).unwrap();
1028 assert_eq!(second_recording.request.url_path, "/second");
1029 assert_eq!(
1030 second_recording.response.body.as_ref().unwrap(),
1031 "second response"
1032 );
1033
1034 mock1.assert_async().await;
1035 mock2.assert_async().await;
1036 }
1037
1038 #[tokio::test]
1039 async fn test_query_parameter_recording() {
1040 let temp_dir = TempDir::new().unwrap();
1042 let temp_path = temp_dir.path().to_path_buf();
1043
1044 let mut server = mockito::Server::new_async().await;
1046
1047 let mock = server
1049 .mock("GET", "/catalogs?max_results=10&page_token=abc123")
1050 .with_status(200)
1051 .with_header("content-type", "application/json")
1052 .with_body(r#"{"catalogs": []}"#)
1053 .create_async()
1054 .await;
1055
1056 let mut client = CloudClient::new_unauthenticated();
1058 client.set_recording_dir(temp_path.clone()).unwrap();
1059
1060 let url = format!("{}/catalogs?max_results=10&page_token=abc123", server.url());
1062 let response = client.get(&url).send().await.unwrap();
1063
1064 assert!(response.status().is_success());
1065
1066 let recordings: Vec<_> = fs::read_dir(&temp_path)
1068 .unwrap()
1069 .filter_map(|entry| {
1070 let entry = entry.ok()?;
1071 let path = entry.path();
1072 if path.extension()? == "json" {
1073 Some(path)
1074 } else {
1075 None
1076 }
1077 })
1078 .collect();
1079
1080 assert_eq!(recordings.len(), 1);
1081
1082 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
1084 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
1085
1086 assert_eq!(recording.request.method, "GET");
1088 assert_eq!(
1089 recording.request.url_path,
1090 "/catalogs?max_results=10&page_token=abc123"
1091 );
1092 assert_eq!(recording.request.body, None);
1093
1094 assert_eq!(recording.response.status, 200);
1096 assert_eq!(
1097 recording.response.body.as_ref().unwrap(),
1098 r#"{"catalogs": []}"#
1099 );
1100
1101 mock.assert_async().await;
1102 }
1103
1104 #[tokio::test]
1109 async fn test_recording_does_not_contain_bearer_token_value() {
1110 let temp_dir = TempDir::new().unwrap();
1111 let temp_path = temp_dir.path().to_path_buf();
1112
1113 let mut server = mockito::Server::new_async().await;
1114 let mock = server
1115 .mock("GET", "/secret")
1116 .match_header("authorization", mockito::Matcher::Any)
1117 .with_status(200)
1118 .with_body("ok")
1120 .create_async()
1121 .await;
1122
1123 let mut client = CloudClient::new_with_token("super-secret-token-12345");
1124 client.set_recording_dir(temp_path.clone()).unwrap();
1125
1126 let url = format!("{}/secret", server.url());
1127 client.get(&url).send().await.unwrap();
1128
1129 let recording_path = temp_path.join("0000.json");
1130 let content = fs::read_to_string(&recording_path).unwrap();
1131
1132 assert!(
1134 !content.contains("super-secret-token-12345"),
1135 "raw bearer token leaked into recording: {content}"
1136 );
1137
1138 mock.assert_async().await;
1139 }
1140
1141 #[tokio::test]
1145 async fn test_recording_redacts_sensitive_response_headers() {
1146 let temp_dir = TempDir::new().unwrap();
1147 let temp_path = temp_dir.path().to_path_buf();
1148
1149 let mut server = mockito::Server::new_async().await;
1150 let mock = server
1152 .mock("GET", "/s3")
1153 .with_status(200)
1154 .with_header("x-amz-security-token", "AQoXnyc4LLI2AJvUAMOGAR8a1234567890")
1155 .with_header("authorization", "Bearer should-be-redacted")
1156 .with_body("{}")
1157 .create_async()
1158 .await;
1159
1160 let mut client = CloudClient::new_unauthenticated();
1161 client.set_recording_dir(temp_path.clone()).unwrap();
1162
1163 let url = format!("{}/s3", server.url());
1164 client.get(&url).send().await.unwrap();
1165
1166 let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1167 assert!(
1168 !content.contains("AQoXnyc4LLI2AJvUAMOGAR8a1234567890"),
1169 "x-amz-security-token leaked into recording: {content}"
1170 );
1171 assert!(
1172 !content.contains("should-be-redacted"),
1173 "Authorization value leaked into recording: {content}"
1174 );
1175 assert_eq!(
1177 content.matches("<REDACTED>").count(),
1178 2,
1179 "expected 2 <REDACTED> entries in recording: {content}"
1180 );
1181
1182 mock.assert_async().await;
1183 }
1184
1185 #[tokio::test]
1187 async fn test_recording_remains_valid_json_after_redaction() {
1188 let temp_dir = TempDir::new().unwrap();
1189 let temp_path = temp_dir.path().to_path_buf();
1190
1191 let mut server = mockito::Server::new_async().await;
1192 let _mock = server
1193 .mock("GET", "/check")
1194 .with_status(200)
1195 .with_header("authorization", "Bearer should-be-redacted")
1196 .with_body(r#"{"ok":true}"#)
1197 .create_async()
1198 .await;
1199
1200 let mut client = CloudClient::new_unauthenticated();
1201 client.set_recording_dir(temp_path.clone()).unwrap();
1202
1203 let url = format!("{}/check", server.url());
1204 client.get(&url).send().await.unwrap();
1205
1206 let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1207 let parsed: RequestResponseInfo = serde_json::from_str(&content)
1209 .expect("recording file must be valid JSON even after redaction");
1210 assert_eq!(parsed.response.status, 200);
1211 }
1212}
1213
1214#[cfg(all(test, not(feature = "recording")))]
1219mod retry_integration_tests {
1220 use super::*;
1221 use std::time::Duration;
1222
1223 fn fast_retry(max_retries: usize) -> RetryConfig {
1224 RetryConfig {
1225 backoff: crate::backoff::BackoffConfig {
1226 init_backoff: Duration::from_millis(1),
1227 max_backoff: Duration::from_millis(5),
1228 base: 2.,
1229 },
1230 max_retries,
1231 retry_timeout: Duration::from_secs(30),
1232 }
1233 }
1234
1235 #[tokio::test]
1239 async fn send_retries_server_error_then_succeeds() {
1240 let mut server = mockito::Server::new_async().await;
1241 let fail = server
1242 .mock("GET", "/r")
1243 .with_status(503)
1244 .expect(1)
1245 .create_async()
1246 .await;
1247 let ok = server
1248 .mock("GET", "/r")
1249 .with_status(200)
1250 .with_body("ok")
1251 .expect(1)
1252 .create_async()
1253 .await;
1254
1255 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1256 let resp = client
1257 .get(format!("{}/r", server.url()))
1258 .send()
1259 .await
1260 .unwrap();
1261
1262 assert_eq!(resp.status(), 200);
1263 assert_eq!(resp.text().await.unwrap(), "ok");
1264 fail.assert_async().await;
1265 ok.assert_async().await;
1266 }
1267
1268 #[tokio::test]
1271 async fn send_does_not_retry_client_error() {
1272 let mut server = mockito::Server::new_async().await;
1273 let mock = server
1274 .mock("GET", "/r")
1275 .with_status(404)
1276 .expect(1)
1277 .create_async()
1278 .await;
1279
1280 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1281 let err = client
1282 .get(format!("{}/r", server.url()))
1283 .send()
1284 .await
1285 .unwrap_err();
1286
1287 assert!(matches!(err, Error::NotFound { .. }), "got {err:?}");
1288 mock.assert_async().await;
1289 }
1290
1291 #[tokio::test]
1296 async fn send_raw_preserves_status_and_body() {
1297 let mut server = mockito::Server::new_async().await;
1298
1299 let not_found = server
1301 .mock("GET", "/nf")
1302 .with_status(404)
1303 .with_body(r#"{"error":{"type":"NoSuchTableException"}}"#)
1304 .expect(1)
1305 .create_async()
1306 .await;
1307 let conflict = server
1309 .mock("POST", "/cf")
1310 .with_status(409)
1311 .with_body(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1312 .expect(1)
1313 .create_async()
1314 .await;
1315 let fail = server
1317 .mock("GET", "/rt")
1318 .with_status(503)
1319 .expect(1)
1320 .create_async()
1321 .await;
1322 let ok = server
1323 .mock("GET", "/rt")
1324 .with_status(200)
1325 .with_body("ok")
1326 .expect(1)
1327 .create_async()
1328 .await;
1329
1330 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1331
1332 let err = client
1333 .get(format!("{}/nf", server.url()))
1334 .send_raw()
1335 .await
1336 .unwrap_err();
1337 match err {
1338 SendRawError::Retry(e) => {
1339 assert_eq!(e.status(), Some(reqwest::StatusCode::NOT_FOUND));
1340 assert_eq!(
1341 e.body(),
1342 Some(r#"{"error":{"type":"NoSuchTableException"}}"#)
1343 );
1344 }
1345 other => panic!("expected Retry error, got {other:?}"),
1346 }
1347
1348 let err = client
1349 .post(format!("{}/cf", server.url()))
1350 .send_raw()
1351 .await
1352 .unwrap_err();
1353 match err {
1354 SendRawError::Retry(e) => {
1355 assert_eq!(e.status(), Some(reqwest::StatusCode::CONFLICT));
1356 assert_eq!(
1357 e.body(),
1358 Some(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1359 );
1360 }
1361 other => panic!("expected Retry error, got {other:?}"),
1362 }
1363
1364 let resp = client
1365 .get(format!("{}/rt", server.url()))
1366 .send_raw()
1367 .await
1368 .unwrap();
1369 assert_eq!(resp.status(), 200);
1370 assert_eq!(resp.text().await.unwrap(), "ok");
1371
1372 not_found.assert_async().await;
1373 conflict.assert_async().await;
1374 fail.assert_async().await;
1375 ok.assert_async().await;
1376 }
1377
1378 #[tokio::test]
1381 async fn sign_and_send_signs_and_retries() {
1382 let mut server = mockito::Server::new_async().await;
1383 let fail = server
1384 .mock("POST", "/rpc")
1385 .match_header("authorization", "Bearer tok")
1386 .with_status(503)
1387 .expect(1)
1388 .create_async()
1389 .await;
1390 let ok = server
1391 .mock("POST", "/rpc")
1392 .match_header("authorization", "Bearer tok")
1393 .with_status(200)
1394 .with_body("pong")
1395 .expect(1)
1396 .create_async()
1397 .await;
1398
1399 let client = CloudClient::new_with_token("tok").with_retry_config(fast_retry(3));
1400 let request = reqwest::Client::new()
1401 .post(format!("{}/rpc", server.url()))
1402 .body("ping")
1403 .build()
1404 .unwrap();
1405 let resp = client.sign_and_send(request).await.unwrap();
1406
1407 assert_eq!(resp.status(), 200);
1408 assert_eq!(resp.text().await.unwrap(), "pong");
1409 fail.assert_async().await;
1410 ok.assert_async().await;
1411 }
1412}
1413
1414#[cfg(test)]
1415mod databricks_client_tests {
1416 use super::*;
1417 use databricks::{DatabricksBuilder, DatabricksConfigKey};
1418
1419 #[tokio::test]
1421 async fn new_databricks_from_builder_signs_with_builder_credential() {
1422 let mut server = mockito::Server::new_async().await;
1423 let mock = server
1424 .mock("GET", "/api/2.1/unity-catalog/catalogs")
1425 .match_header("authorization", "Bearer builder-pat")
1426 .with_status(200)
1427 .create_async()
1428 .await;
1429
1430 let builder = DatabricksBuilder::new()
1431 .with_config(DatabricksConfigKey::Host, server.url())
1432 .with_config(DatabricksConfigKey::Token, "builder-pat")
1433 .with_config(
1434 DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1435 "true",
1436 );
1437 let client = CloudClient::new_databricks_from_builder(builder, None).unwrap();
1438
1439 let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1440 let response = client.get(&url).send().await.unwrap();
1441 assert_eq!(response.status(), 200);
1442 mock.assert_async().await;
1443 }
1444
1445 #[tokio::test]
1447 async fn new_databricks_applies_string_options() {
1448 let mut server = mockito::Server::new_async().await;
1449 let mock = server
1450 .mock("GET", "/ping")
1451 .match_header("authorization", "Bearer option-pat")
1452 .with_status(200)
1453 .create_async()
1454 .await;
1455
1456 let client = CloudClient::new_databricks(
1457 [
1458 ("databricks_host", server.url()),
1459 ("databricks_token", "option-pat".to_string()),
1460 ("allow_http", "true".to_string()),
1461 ("not_a_real_key", "ignored".to_string()),
1462 ],
1463 None,
1464 )
1465 .unwrap();
1466
1467 let response = client
1468 .get(format!("{}/ping", server.url()))
1469 .send()
1470 .await
1471 .unwrap();
1472 assert_eq!(response.status(), 200);
1473 mock.assert_async().await;
1474 }
1475}
1476
1477#[cfg(test)]
1478mod databricks_azure_auth_tests {
1479 use super::*;
1480 use databricks::{DatabricksBuilder, DatabricksConfigKey};
1481
1482 fn azure_sp_builder(server: &mockito::Server) -> DatabricksBuilder {
1483 DatabricksBuilder::new()
1484 .with_config(DatabricksConfigKey::Host, server.url())
1485 .with_config(DatabricksConfigKey::ClientId, "sp-id")
1486 .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1487 .with_config(DatabricksConfigKey::AzureTenantId, "tenant-1")
1488 .with_config(DatabricksConfigKey::AzureAuthorityHost, server.url())
1489 .with_config(
1490 DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1491 "true",
1492 )
1493 }
1494
1495 #[tokio::test]
1498 async fn tenant_selects_azure_client_secret_auth() {
1499 let mut server = mockito::Server::new_async().await;
1500 let aad = server
1501 .mock("POST", "/tenant-1/oauth2/v2.0/token")
1502 .match_body(mockito::Matcher::AllOf(vec![
1503 mockito::Matcher::UrlEncoded("client_id".into(), "sp-id".into()),
1504 mockito::Matcher::UrlEncoded(
1505 "scope".into(),
1506 "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default".into(),
1507 ),
1508 ]))
1509 .with_status(200)
1510 .with_header("content-type", "application/json")
1511 .with_body(r#"{"access_token":"aad-token","expires_in":3600,"token_type":"Bearer"}"#)
1512 .expect(1)
1513 .create_async()
1514 .await;
1515 let m2m = server
1516 .mock("POST", "/oidc/v1/token")
1517 .expect(0)
1518 .create_async()
1519 .await;
1520 let api = server
1521 .mock("GET", "/api/2.1/unity-catalog/catalogs")
1522 .match_header("authorization", "Bearer aad-token")
1523 .with_status(200)
1524 .create_async()
1525 .await;
1526
1527 let client =
1528 CloudClient::new_databricks_from_builder(azure_sp_builder(&server), None).unwrap();
1529 let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1530 assert_eq!(client.get(&url).send().await.unwrap().status(), 200);
1531
1532 aad.assert_async().await;
1533 m2m.assert_async().await;
1534 api.assert_async().await;
1535 }
1536
1537 #[test]
1538 fn forced_azure_client_secret_requires_a_tenant() {
1539 let result = DatabricksBuilder::new()
1540 .with_config(DatabricksConfigKey::ClientId, "sp-id")
1541 .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1542 .with_auth_type("azure-client-secret")
1543 .build(None);
1544 let Err(err) = result else {
1545 panic!("expected an error without a tenant id")
1546 };
1547 assert!(err.to_string().contains("AZURE_TENANT_ID"), "{err}");
1548 }
1549
1550 #[test]
1551 fn azure_config_keys_roundtrip() {
1552 for key in [
1553 DatabricksConfigKey::AzureTenantId,
1554 DatabricksConfigKey::AzureAuthorityHost,
1555 ] {
1556 assert_eq!(key.as_ref().parse::<DatabricksConfigKey>().unwrap(), key);
1557 }
1558 assert_eq!(
1559 "azure_tenant_id".parse::<DatabricksConfigKey>().unwrap(),
1560 DatabricksConfigKey::AzureTenantId
1561 );
1562 }
1563}