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;
72pub mod bearer;
73mod client;
74mod config;
75#[cfg(feature = "connectrpc")]
76pub mod connectrpc;
77mod credential;
78pub mod databricks;
79mod error;
80pub mod gcp;
81pub mod service;
82
83mod retry;
84mod token;
85mod util;
86
87pub use client::{Certificate, ClientConfigKey, ClientOptions};
88pub use credential::*;
89pub use error::*;
90pub use retry::Error as RetryError;
91pub use retry::RetryConfig;
92pub use service::{ReqwestService, SpawnService};
93
94type SharedSigner = Arc<dyn RequestSigner>;
96
97#[derive(Debug)]
99struct NoopSigner;
100
101impl RequestSigner for NoopSigner {
102 fn sign<'a>(
103 &'a self,
104 req: RequestBuilder,
105 ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
106 Box::pin(async move { Ok(req) })
107 }
108}
109
110#[derive(Debug)]
112struct BearerTokenSigner {
113 token: String,
114}
115
116impl RequestSigner for BearerTokenSigner {
117 fn sign<'a>(
118 &'a self,
119 req: RequestBuilder,
120 ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
121 let token = self.token.clone();
122 Box::pin(async move { Ok(req.bearer_auth(&token)) })
123 }
124}
125
126#[cfg(feature = "recording")]
127#[derive(Debug, Clone)]
128struct RecordingState {
129 out_dir: PathBuf,
130 counter: Arc<AtomicU64>,
131}
132
133#[derive(Clone)]
140pub struct CloudClient {
141 signer: SharedSigner,
142 reqwest_client: Client,
143 service: Arc<dyn HttpService>,
144 pub retry_config: RetryConfig,
151 #[cfg(feature = "recording")]
152 recording: Option<RecordingState>,
153}
154
155impl CloudClient {
156 fn new_with_signer(
157 signer: SharedSigner,
158 reqwest_client: Client,
159 service: Arc<dyn HttpService>,
160 retry_config: RetryConfig,
161 ) -> Self {
162 Self {
163 signer,
164 reqwest_client,
165 service,
166 retry_config,
167 #[cfg(feature = "recording")]
168 recording: None,
169 }
170 }
171
172 pub fn new_aws<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
177 where
178 I: IntoIterator<Item = (K, V)>,
179 K: AsRef<str>,
180 V: Into<String>,
181 {
182 let config = options
183 .into_iter()
184 .fold(
185 aws::AmazonBuilder::new(),
186 |builder, (key, value)| match key.as_ref().parse() {
187 Ok(k) => builder.with_config(k, value),
188 Err(_) => builder,
189 },
190 )
191 .build(runtime)?;
192
193 let reqwest_client = config.client_options.client()?;
194 let service = make_service(reqwest_client.clone(), runtime);
195 let retry_config = config.retry_config.clone();
196 Ok(Self::new_with_signer(
197 Arc::new(config),
198 reqwest_client,
199 service,
200 retry_config,
201 ))
202 }
203
204 pub fn new_google<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
209 where
210 I: IntoIterator<Item = (K, V)>,
211 K: AsRef<str>,
212 V: Into<String>,
213 {
214 let config = options
215 .into_iter()
216 .fold(
217 gcp::GoogleBuilder::new(),
218 |builder, (key, value)| match key.as_ref().parse() {
219 Ok(k) => builder.with_config(k, value),
220 Err(_) => builder,
221 },
222 )
223 .build(runtime)?;
224
225 let reqwest_client = config.client_options.client()?;
226 let service = make_service(reqwest_client.clone(), runtime);
227 let retry_config = config.retry_config.clone();
228 Ok(Self::new_with_signer(
229 Arc::new(config),
230 reqwest_client,
231 service,
232 retry_config,
233 ))
234 }
235
236 pub fn new_azure<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
241 where
242 I: IntoIterator<Item = (K, V)>,
243 K: AsRef<str>,
244 V: Into<String>,
245 {
246 let config = options
247 .into_iter()
248 .fold(
249 azure::AzureBuilder::new(),
250 |builder, (key, value)| match key.as_ref().parse() {
251 Ok(k) => builder.with_config(k, value),
252 Err(_) => builder,
253 },
254 )
255 .build(runtime)?;
256
257 let reqwest_client = config.client_options.client()?;
258 let service = make_service(reqwest_client.clone(), runtime);
259 let retry_config = config.retry_config.clone();
260 Ok(Self::new_with_signer(
261 Arc::new(config),
262 reqwest_client,
263 service,
264 retry_config,
265 ))
266 }
267
268 pub fn new_databricks<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
280 where
281 I: IntoIterator<Item = (K, V)>,
282 K: AsRef<str>,
283 V: Into<String>,
284 {
285 let builder = options.into_iter().fold(
286 databricks::DatabricksBuilder::new(),
287 |builder, (key, value)| match key.as_ref().parse() {
288 Ok(k) => builder.with_config(k, value),
289 Err(_) => builder,
290 },
291 );
292 Self::new_databricks_from_builder(builder, runtime)
293 }
294
295 pub fn new_databricks_from_builder(
306 builder: databricks::DatabricksBuilder,
307 runtime: Option<&Handle>,
308 ) -> Result<Self> {
309 let config = builder.build(runtime)?;
310
311 let reqwest_client = config.client_options.client()?;
312 let service = make_service(reqwest_client.clone(), runtime);
313 let retry_config = config.retry_config.clone();
314 Ok(Self::new_with_signer(
315 Arc::new(config),
316 reqwest_client,
317 service,
318 retry_config,
319 ))
320 }
321
322 pub fn new_with_bearer(
328 provider: bearer::BearerCredentialProvider,
329 options: &ClientOptions,
330 retry: RetryConfig,
331 runtime: Option<&Handle>,
332 ) -> Result<Self> {
333 let reqwest_client = options.client()?;
334 let service = make_service(reqwest_client.clone(), runtime);
335 Ok(Self::new_with_signer(
336 Arc::new(bearer::BearerProviderSigner(provider)),
337 reqwest_client,
338 service,
339 retry,
340 ))
341 }
342
343 pub fn new_with_token(token: impl ToString) -> Self {
345 let reqwest_client = Client::new();
346 let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
347 Self::new_with_signer(
348 Arc::new(BearerTokenSigner {
349 token: token.to_string(),
350 }),
351 reqwest_client,
352 service,
353 RetryConfig::default(),
354 )
355 }
356
357 pub fn new_unauthenticated() -> Self {
359 let reqwest_client = Client::new();
360 let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
361 Self::new_with_signer(
362 Arc::new(NoopSigner),
363 reqwest_client,
364 service,
365 RetryConfig::default(),
366 )
367 }
368
369 pub fn new_unauthenticated_with_options(
375 options: &ClientOptions,
376 retry: RetryConfig,
377 runtime: Option<&Handle>,
378 ) -> Result<Self> {
379 let reqwest_client = options.client()?;
380 let service = make_service(reqwest_client.clone(), runtime);
381 Ok(Self::new_with_signer(
382 Arc::new(NoopSigner),
383 reqwest_client,
384 service,
385 retry,
386 ))
387 }
388
389 pub fn with_runtime(mut self, handle: Handle) -> Self {
395 self.service = Arc::new(SpawnService::new(self.service, handle));
396 self
397 }
398
399 pub fn with_http_service(mut self, service: Arc<dyn HttpService>) -> Self {
401 self.service = service;
402 self
403 }
404
405 pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
413 self.retry_config = retry_config;
414 self
415 }
416
417 pub async fn sign_and_send(&self, request: reqwest::Request) -> Result<reqwest::Response> {
435 let builder = self
439 .reqwest_client
440 .request(request.method().clone(), request.url().clone());
441 let builder = builder.headers(request.headers().clone());
442 let builder = match request.body().and_then(|b| b.as_bytes()) {
443 Some(bytes) => builder.body(bytes.to_vec()),
444 None => builder,
445 };
446 let builder = self.signer.sign(builder).await?;
447 builder
448 .send_retry(&self.retry_config, self.service.clone())
449 .await
450 .map_err(|e| e.error())
451 }
452
453 pub fn request<U: IntoUrl>(&self, method: Method, url: U) -> CloudRequestBuilder {
459 CloudRequestBuilder {
460 builder: self.reqwest_client.request(method, url),
461 client: self.clone(),
462 #[cfg(feature = "recording")]
463 out_dir: self.recording.as_ref().map(|r| r.out_dir.clone()),
464 }
465 }
466
467 pub fn get<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
469 self.request(Method::GET, url)
470 }
471
472 pub fn post<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
474 self.request(Method::POST, url)
475 }
476
477 pub fn put<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
479 self.request(Method::PUT, url)
480 }
481
482 pub fn delete<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
484 self.request(Method::DELETE, url)
485 }
486
487 pub fn head<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
489 self.request(Method::HEAD, url)
490 }
491
492 pub fn patch<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
494 self.request(Method::PATCH, url)
495 }
496
497 pub fn options<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
499 self.request(Method::OPTIONS, url)
500 }
501
502 pub fn trace<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
504 self.request(Method::TRACE, url)
505 }
506
507 pub fn connect<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
509 self.request(Method::CONNECT, url)
510 }
511
512 #[cfg(feature = "recording")]
531 pub fn set_recording_dir(&mut self, out_dir: std::path::PathBuf) -> Result<(), std::io::Error> {
532 let out_dir = std::fs::canonicalize(out_dir)?;
533 self.recording = Some(RecordingState {
534 out_dir,
535 counter: Arc::new(AtomicU64::new(0)),
536 });
537 Ok(())
538 }
539}
540
541pub struct CloudRequestBuilder {
547 builder: RequestBuilder,
548 client: CloudClient,
549 #[cfg(feature = "recording")]
550 out_dir: Option<PathBuf>,
551}
552
553impl CloudRequestBuilder {
554 pub fn header<K, V>(mut self, key: K, value: V) -> CloudRequestBuilder
556 where
557 HeaderName: TryFrom<K>,
558 <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
559 HeaderValue: TryFrom<V>,
560 <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
561 {
562 self.builder = self.builder.header(key, value);
563 self
564 }
565
566 pub fn headers(mut self, headers: HeaderMap) -> CloudRequestBuilder {
570 self.builder = self.builder.headers(headers);
571 self
572 }
573
574 pub fn body<T: Into<Body>>(mut self, body: T) -> CloudRequestBuilder {
576 self.builder = self.builder.body(body);
577 self
578 }
579
580 pub fn timeout(mut self, timeout: Duration) -> CloudRequestBuilder {
586 self.builder = self.builder.timeout(timeout);
587 self
588 }
589
590 pub fn query<T: Serialize + ?Sized>(mut self, query: &T) -> CloudRequestBuilder {
609 self.builder = self.builder.query(query);
610 self
611 }
612
613 pub fn json<T: Serialize + ?Sized>(mut self, json: &T) -> CloudRequestBuilder {
620 self.builder = self.builder.json(json);
621 self
622 }
623
624 pub async fn send(mut self) -> Result<reqwest::Response> {
647 self.builder = self.client.signer.sign(self.builder).await?;
648
649 #[cfg(not(feature = "recording"))]
650 {
651 self.builder
652 .send_retry(&self.client.retry_config, self.client.service.clone())
653 .await
654 .map_err(|e| e.error())
655 }
656 #[cfg(feature = "recording")]
657 {
658 let response = send_record(self).await?;
659 Ok(response)
660 }
661 }
662
663 #[cfg(not(feature = "recording"))]
690 pub async fn send_raw(mut self) -> Result<reqwest::Response, SendRawError> {
691 self.builder = self.client.signer.sign(self.builder).await?;
692 self.builder
693 .send_retry(&self.client.retry_config, self.client.service.clone())
694 .await
695 .map_err(SendRawError::Retry)
696 }
697
698 #[cfg(feature = "recording")]
701 pub async fn send_raw(self) -> Result<reqwest::Response, SendRawError> {
702 send_record(self).await
703 }
704}
705
706#[cfg(not(feature = "recording"))]
714#[derive(Debug, thiserror::Error)]
715pub enum SendRawError {
716 #[error(transparent)]
718 Sign(#[from] Error),
719 #[error(transparent)]
721 Retry(RetryError),
722}
723
724#[cfg(feature = "recording")]
727pub type SendRawError = Error;
728
729#[cfg(feature = "recording")]
730#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
731pub struct RequestResponseInfo {
732 pub request: RequestInfo,
733 pub response: ResponseInfo,
734}
735
736#[cfg(feature = "recording")]
737#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
738pub struct RequestInfo {
739 pub method: String,
740 pub url_path: String,
741 pub body: Option<String>,
742}
743
744#[cfg(feature = "recording")]
745#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
746pub struct ResponseInfo {
747 pub status: u16,
748 pub headers: HashMap<String, String>,
749 pub body: Option<String>,
750}
751
752#[cfg(feature = "recording")]
757const SENSITIVE_HEADERS: &[&str] = &[
758 "authorization",
759 "x-amz-security-token",
760 "x-amz-content-sha256",
761 "x-databricks-authorization",
762 "x-ms-identity-principal-id",
763 "x-goog-iam-credentials-token",
764 "cookie",
765 "set-cookie",
766];
767
768#[cfg(feature = "recording")]
769fn redact_headers(headers: &HashMap<String, String>) -> HashMap<String, String> {
770 headers
771 .iter()
772 .map(|(k, v)| {
773 let v = if SENSITIVE_HEADERS.contains(&k.to_lowercase().as_str()) {
774 "<REDACTED>".to_string()
775 } else {
776 v.clone()
777 };
778 (k.clone(), v)
779 })
780 .collect()
781}
782
783#[cfg(feature = "recording")]
784async fn send_record(builder: CloudRequestBuilder) -> Result<reqwest::Response> {
785 let Some(out_dir) = builder.out_dir else {
786 let request = builder.builder.build().expect("request to be valid");
787 return builder.client.service.call(request).await;
788 };
789 let (_client, request) = builder.builder.build_split();
790 let request = request.expect("request to be valid");
791
792 let request_info = RequestInfo {
793 method: request.method().as_str().to_string(),
794 url_path: {
795 let url = request.url();
796 match url.query() {
797 Some(query) => format!("{}?{}", url.path(), query),
798 None => url.path().to_string(),
799 }
800 },
801 body: request
802 .body()
803 .and_then(|b| b.as_bytes().map(|b| String::from_utf8_lossy(b).to_string())),
804 };
805
806 let response = builder.client.service.call(request).await?;
807
808 let status = response.status().as_u16();
810 let raw_headers: HashMap<String, String> = response
811 .headers()
812 .iter()
813 .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
814 .collect();
815
816 let response_bytes = response.bytes().await?;
818 let response_body = if response_bytes.is_empty() {
819 None
820 } else {
821 Some(String::from_utf8_lossy(&response_bytes).to_string())
822 };
823
824 let recording = RequestResponseInfo {
825 request: request_info,
826 response: ResponseInfo {
827 status,
828 headers: redact_headers(&raw_headers),
830 body: response_body,
831 },
832 };
833
834 let counter = builder
835 .client
836 .recording
837 .as_ref()
838 .map(|r| r.counter.fetch_add(1, Ordering::SeqCst))
839 .unwrap_or(0);
840 let file_path = out_dir.join(format!("{counter:04}.json"));
841 if let Err(e) = std::fs::File::create(&file_path)
842 .and_then(|f| serde_json::to_writer_pretty(f, &recording).map_err(Into::into))
843 {
844 tracing::warn!(
845 "Failed to write recording to {}: {}",
846 file_path.display(),
847 e
848 );
849 }
850
851 let mut mock_response = http::Response::builder().status(status);
853 for (k, v) in &raw_headers {
854 mock_response = mock_response.header(k, v);
855 }
856 let mock_response = mock_response
857 .body(response_bytes)
858 .expect("valid status code and headers");
859
860 Ok(reqwest::Response::from(mock_response))
861}
862
863#[cfg(all(test, feature = "recording"))]
864mod tests {
865 use super::*;
866 use std::fs;
867 use tempfile::TempDir;
868
869 #[tokio::test]
870 async fn test_request_response_recording() {
871 let temp_dir = TempDir::new().unwrap();
873 let temp_path = temp_dir.path().to_path_buf();
874
875 let mut server = mockito::Server::new_async().await;
877 let mock = server
878 .mock("GET", "/test")
879 .with_status(200)
880 .with_header("content-type", "application/json")
881 .with_body(r#"{"message": "Hello, World!"}"#)
882 .create_async()
883 .await;
884
885 let mut client = CloudClient::new_unauthenticated();
887 client.set_recording_dir(temp_path.clone()).unwrap();
888
889 let url = format!("{}/test", server.url());
891 let response = client.get(&url).send().await.unwrap();
892
893 assert_eq!(response.status(), 200);
895 let body = response.text().await.unwrap();
896 assert_eq!(body, r#"{"message": "Hello, World!"}"#);
897
898 let recordings: Vec<_> = fs::read_dir(&temp_path)
900 .unwrap()
901 .filter_map(|entry| {
902 let entry = entry.ok()?;
903 let path = entry.path();
904 if path.extension()? == "json" {
905 Some(path)
906 } else {
907 None
908 }
909 })
910 .collect();
911
912 assert_eq!(recordings.len(), 1, "Expected exactly one recording file");
913
914 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
916 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
917
918 assert_eq!(recording.request.method, "GET");
920 assert_eq!(recording.request.url_path, "/test");
921 assert_eq!(recording.request.body, None);
922
923 assert_eq!(recording.response.status, 200);
925 assert_eq!(
926 recording.response.headers.get("content-type").unwrap(),
927 "application/json"
928 );
929 assert_eq!(
930 recording.response.body.as_ref().unwrap(),
931 r#"{"message": "Hello, World!"}"#
932 );
933
934 mock.assert_async().await;
935 }
936
937 #[tokio::test]
938 async fn test_recording_with_request_body() {
939 let temp_dir = TempDir::new().unwrap();
941 let temp_path = temp_dir.path().to_path_buf();
942
943 let mut server = mockito::Server::new_async().await;
945 let mock = server
946 .mock("POST", "/create")
947 .with_status(201)
948 .with_header("location", "/resource/123")
949 .with_body(r#"{"id": 123, "status": "created"}"#)
950 .create_async()
951 .await;
952
953 let mut client = CloudClient::new_unauthenticated();
955 client.set_recording_dir(temp_path.clone()).unwrap();
956
957 let url = format!("{}/create", server.url());
959 let response = client
960 .post(&url)
961 .json(&serde_json::json!({"name": "test resource"}))
962 .send()
963 .await
964 .unwrap();
965
966 assert_eq!(response.status(), 201);
968
969 let recordings: Vec<_> = fs::read_dir(&temp_path)
971 .unwrap()
972 .filter_map(|entry| {
973 let entry = entry.ok()?;
974 let path = entry.path();
975 if path.extension()? == "json" {
976 Some(path)
977 } else {
978 None
979 }
980 })
981 .collect();
982
983 assert_eq!(recordings.len(), 1);
984
985 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
987 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
988
989 assert_eq!(recording.request.method, "POST");
991 assert_eq!(recording.request.url_path, "/create");
992 assert!(recording.request.body.is_some());
993 assert!(recording.request.body.unwrap().contains("test resource"));
994
995 assert_eq!(recording.response.status, 201);
997 assert_eq!(
998 recording.response.headers.get("location").unwrap(),
999 "/resource/123"
1000 );
1001 assert!(recording.response.body.unwrap().contains("created"));
1002
1003 mock.assert_async().await;
1004 }
1005
1006 #[tokio::test]
1007 async fn test_counter_based_file_naming() {
1008 let temp_dir = TempDir::new().unwrap();
1010 let temp_path = temp_dir.path().to_path_buf();
1011
1012 let mut server = mockito::Server::new_async().await;
1014 let mock1 = server
1015 .mock("GET", "/first")
1016 .with_status(200)
1017 .with_body("first response")
1018 .create_async()
1019 .await;
1020 let mock2 = server
1021 .mock("GET", "/second")
1022 .with_status(200)
1023 .with_body("second response")
1024 .create_async()
1025 .await;
1026
1027 let mut client = CloudClient::new_unauthenticated();
1029 client.set_recording_dir(temp_path.clone()).unwrap();
1030
1031 let url1 = format!("{}/first", server.url());
1033 let url2 = format!("{}/second", server.url());
1034
1035 let _response1 = client.get(&url1).send().await.unwrap();
1036 let _response2 = client.get(&url2).send().await.unwrap();
1037
1038 let mut recordings: Vec<_> = fs::read_dir(&temp_path)
1040 .unwrap()
1041 .filter_map(|entry| {
1042 let entry = entry.ok()?;
1043 let path = entry.path();
1044 if path.extension()? == "json" {
1045 Some(path)
1046 } else {
1047 None
1048 }
1049 })
1050 .collect();
1051
1052 recordings.sort();
1053 assert_eq!(recordings.len(), 2);
1054
1055 assert!(recordings[0].file_name().unwrap().to_str().unwrap() == "0000.json");
1057 assert!(recordings[1].file_name().unwrap().to_str().unwrap() == "0001.json");
1058
1059 let first_content = fs::read_to_string(&recordings[0]).unwrap();
1061 let first_recording: RequestResponseInfo = serde_json::from_str(&first_content).unwrap();
1062 assert_eq!(first_recording.request.url_path, "/first");
1063 assert_eq!(
1064 first_recording.response.body.as_ref().unwrap(),
1065 "first response"
1066 );
1067
1068 let second_content = fs::read_to_string(&recordings[1]).unwrap();
1069 let second_recording: RequestResponseInfo = serde_json::from_str(&second_content).unwrap();
1070 assert_eq!(second_recording.request.url_path, "/second");
1071 assert_eq!(
1072 second_recording.response.body.as_ref().unwrap(),
1073 "second response"
1074 );
1075
1076 mock1.assert_async().await;
1077 mock2.assert_async().await;
1078 }
1079
1080 #[tokio::test]
1081 async fn test_query_parameter_recording() {
1082 let temp_dir = TempDir::new().unwrap();
1084 let temp_path = temp_dir.path().to_path_buf();
1085
1086 let mut server = mockito::Server::new_async().await;
1088
1089 let mock = server
1091 .mock("GET", "/catalogs?max_results=10&page_token=abc123")
1092 .with_status(200)
1093 .with_header("content-type", "application/json")
1094 .with_body(r#"{"catalogs": []}"#)
1095 .create_async()
1096 .await;
1097
1098 let mut client = CloudClient::new_unauthenticated();
1100 client.set_recording_dir(temp_path.clone()).unwrap();
1101
1102 let url = format!("{}/catalogs?max_results=10&page_token=abc123", server.url());
1104 let response = client.get(&url).send().await.unwrap();
1105
1106 assert!(response.status().is_success());
1107
1108 let recordings: Vec<_> = fs::read_dir(&temp_path)
1110 .unwrap()
1111 .filter_map(|entry| {
1112 let entry = entry.ok()?;
1113 let path = entry.path();
1114 if path.extension()? == "json" {
1115 Some(path)
1116 } else {
1117 None
1118 }
1119 })
1120 .collect();
1121
1122 assert_eq!(recordings.len(), 1);
1123
1124 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
1126 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
1127
1128 assert_eq!(recording.request.method, "GET");
1130 assert_eq!(
1131 recording.request.url_path,
1132 "/catalogs?max_results=10&page_token=abc123"
1133 );
1134 assert_eq!(recording.request.body, None);
1135
1136 assert_eq!(recording.response.status, 200);
1138 assert_eq!(
1139 recording.response.body.as_ref().unwrap(),
1140 r#"{"catalogs": []}"#
1141 );
1142
1143 mock.assert_async().await;
1144 }
1145
1146 #[tokio::test]
1151 async fn test_recording_does_not_contain_bearer_token_value() {
1152 let temp_dir = TempDir::new().unwrap();
1153 let temp_path = temp_dir.path().to_path_buf();
1154
1155 let mut server = mockito::Server::new_async().await;
1156 let mock = server
1157 .mock("GET", "/secret")
1158 .match_header("authorization", mockito::Matcher::Any)
1159 .with_status(200)
1160 .with_body("ok")
1162 .create_async()
1163 .await;
1164
1165 let mut client = CloudClient::new_with_token("super-secret-token-12345");
1166 client.set_recording_dir(temp_path.clone()).unwrap();
1167
1168 let url = format!("{}/secret", server.url());
1169 client.get(&url).send().await.unwrap();
1170
1171 let recording_path = temp_path.join("0000.json");
1172 let content = fs::read_to_string(&recording_path).unwrap();
1173
1174 assert!(
1176 !content.contains("super-secret-token-12345"),
1177 "raw bearer token leaked into recording: {content}"
1178 );
1179
1180 mock.assert_async().await;
1181 }
1182
1183 #[tokio::test]
1187 async fn test_recording_redacts_sensitive_response_headers() {
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
1194 .mock("GET", "/s3")
1195 .with_status(200)
1196 .with_header("x-amz-security-token", "AQoXnyc4LLI2AJvUAMOGAR8a1234567890")
1197 .with_header("authorization", "Bearer should-be-redacted")
1198 .with_body("{}")
1199 .create_async()
1200 .await;
1201
1202 let mut client = CloudClient::new_unauthenticated();
1203 client.set_recording_dir(temp_path.clone()).unwrap();
1204
1205 let url = format!("{}/s3", server.url());
1206 client.get(&url).send().await.unwrap();
1207
1208 let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1209 assert!(
1210 !content.contains("AQoXnyc4LLI2AJvUAMOGAR8a1234567890"),
1211 "x-amz-security-token leaked into recording: {content}"
1212 );
1213 assert!(
1214 !content.contains("should-be-redacted"),
1215 "Authorization value leaked into recording: {content}"
1216 );
1217 assert_eq!(
1219 content.matches("<REDACTED>").count(),
1220 2,
1221 "expected 2 <REDACTED> entries in recording: {content}"
1222 );
1223
1224 mock.assert_async().await;
1225 }
1226
1227 #[tokio::test]
1229 async fn test_recording_remains_valid_json_after_redaction() {
1230 let temp_dir = TempDir::new().unwrap();
1231 let temp_path = temp_dir.path().to_path_buf();
1232
1233 let mut server = mockito::Server::new_async().await;
1234 let _mock = server
1235 .mock("GET", "/check")
1236 .with_status(200)
1237 .with_header("authorization", "Bearer should-be-redacted")
1238 .with_body(r#"{"ok":true}"#)
1239 .create_async()
1240 .await;
1241
1242 let mut client = CloudClient::new_unauthenticated();
1243 client.set_recording_dir(temp_path.clone()).unwrap();
1244
1245 let url = format!("{}/check", server.url());
1246 client.get(&url).send().await.unwrap();
1247
1248 let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1249 let parsed: RequestResponseInfo = serde_json::from_str(&content)
1251 .expect("recording file must be valid JSON even after redaction");
1252 assert_eq!(parsed.response.status, 200);
1253 }
1254}
1255
1256#[cfg(all(test, not(feature = "recording")))]
1261mod retry_integration_tests {
1262 use super::*;
1263 use std::time::Duration;
1264
1265 fn fast_retry(max_retries: usize) -> RetryConfig {
1266 RetryConfig {
1267 backoff: crate::backoff::BackoffConfig {
1268 init_backoff: Duration::from_millis(1),
1269 max_backoff: Duration::from_millis(5),
1270 base: 2.,
1271 },
1272 max_retries,
1273 retry_timeout: Duration::from_secs(30),
1274 }
1275 }
1276
1277 #[tokio::test]
1281 async fn send_retries_server_error_then_succeeds() {
1282 let mut server = mockito::Server::new_async().await;
1283 let fail = server
1284 .mock("GET", "/r")
1285 .with_status(503)
1286 .expect(1)
1287 .create_async()
1288 .await;
1289 let ok = server
1290 .mock("GET", "/r")
1291 .with_status(200)
1292 .with_body("ok")
1293 .expect(1)
1294 .create_async()
1295 .await;
1296
1297 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1298 let resp = client
1299 .get(format!("{}/r", server.url()))
1300 .send()
1301 .await
1302 .unwrap();
1303
1304 assert_eq!(resp.status(), 200);
1305 assert_eq!(resp.text().await.unwrap(), "ok");
1306 fail.assert_async().await;
1307 ok.assert_async().await;
1308 }
1309
1310 #[tokio::test]
1313 async fn send_does_not_retry_client_error() {
1314 let mut server = mockito::Server::new_async().await;
1315 let mock = server
1316 .mock("GET", "/r")
1317 .with_status(404)
1318 .expect(1)
1319 .create_async()
1320 .await;
1321
1322 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1323 let err = client
1324 .get(format!("{}/r", server.url()))
1325 .send()
1326 .await
1327 .unwrap_err();
1328
1329 assert!(matches!(err, Error::NotFound { .. }), "got {err:?}");
1330 mock.assert_async().await;
1331 }
1332
1333 #[tokio::test]
1338 async fn send_raw_preserves_status_and_body() {
1339 let mut server = mockito::Server::new_async().await;
1340
1341 let not_found = server
1343 .mock("GET", "/nf")
1344 .with_status(404)
1345 .with_body(r#"{"error":{"type":"NoSuchTableException"}}"#)
1346 .expect(1)
1347 .create_async()
1348 .await;
1349 let conflict = server
1351 .mock("POST", "/cf")
1352 .with_status(409)
1353 .with_body(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1354 .expect(1)
1355 .create_async()
1356 .await;
1357 let fail = server
1359 .mock("GET", "/rt")
1360 .with_status(503)
1361 .expect(1)
1362 .create_async()
1363 .await;
1364 let ok = server
1365 .mock("GET", "/rt")
1366 .with_status(200)
1367 .with_body("ok")
1368 .expect(1)
1369 .create_async()
1370 .await;
1371
1372 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1373
1374 let err = client
1375 .get(format!("{}/nf", server.url()))
1376 .send_raw()
1377 .await
1378 .unwrap_err();
1379 match err {
1380 SendRawError::Retry(e) => {
1381 assert_eq!(e.status(), Some(reqwest::StatusCode::NOT_FOUND));
1382 assert_eq!(
1383 e.body(),
1384 Some(r#"{"error":{"type":"NoSuchTableException"}}"#)
1385 );
1386 }
1387 other => panic!("expected Retry error, got {other:?}"),
1388 }
1389
1390 let err = client
1391 .post(format!("{}/cf", server.url()))
1392 .send_raw()
1393 .await
1394 .unwrap_err();
1395 match err {
1396 SendRawError::Retry(e) => {
1397 assert_eq!(e.status(), Some(reqwest::StatusCode::CONFLICT));
1398 assert_eq!(
1399 e.body(),
1400 Some(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1401 );
1402 }
1403 other => panic!("expected Retry error, got {other:?}"),
1404 }
1405
1406 let resp = client
1407 .get(format!("{}/rt", server.url()))
1408 .send_raw()
1409 .await
1410 .unwrap();
1411 assert_eq!(resp.status(), 200);
1412 assert_eq!(resp.text().await.unwrap(), "ok");
1413
1414 not_found.assert_async().await;
1415 conflict.assert_async().await;
1416 fail.assert_async().await;
1417 ok.assert_async().await;
1418 }
1419
1420 #[tokio::test]
1423 async fn sign_and_send_signs_and_retries() {
1424 let mut server = mockito::Server::new_async().await;
1425 let fail = server
1426 .mock("POST", "/rpc")
1427 .match_header("authorization", "Bearer tok")
1428 .with_status(503)
1429 .expect(1)
1430 .create_async()
1431 .await;
1432 let ok = server
1433 .mock("POST", "/rpc")
1434 .match_header("authorization", "Bearer tok")
1435 .with_status(200)
1436 .with_body("pong")
1437 .expect(1)
1438 .create_async()
1439 .await;
1440
1441 let client = CloudClient::new_with_token("tok").with_retry_config(fast_retry(3));
1442 let request = reqwest::Client::new()
1443 .post(format!("{}/rpc", server.url()))
1444 .body("ping")
1445 .build()
1446 .unwrap();
1447 let resp = client.sign_and_send(request).await.unwrap();
1448
1449 assert_eq!(resp.status(), 200);
1450 assert_eq!(resp.text().await.unwrap(), "pong");
1451 fail.assert_async().await;
1452 ok.assert_async().await;
1453 }
1454}
1455
1456#[cfg(test)]
1457mod databricks_client_tests {
1458 use super::*;
1459 use databricks::{DatabricksBuilder, DatabricksConfigKey};
1460
1461 #[tokio::test]
1463 async fn new_databricks_from_builder_signs_with_builder_credential() {
1464 let mut server = mockito::Server::new_async().await;
1465 let mock = server
1466 .mock("GET", "/api/2.1/unity-catalog/catalogs")
1467 .match_header("authorization", "Bearer builder-pat")
1468 .with_status(200)
1469 .create_async()
1470 .await;
1471
1472 let builder = DatabricksBuilder::new()
1473 .with_config(DatabricksConfigKey::Host, server.url())
1474 .with_config(DatabricksConfigKey::Token, "builder-pat")
1475 .with_config(
1476 DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1477 "true",
1478 );
1479 let client = CloudClient::new_databricks_from_builder(builder, None).unwrap();
1480
1481 let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1482 let response = client.get(&url).send().await.unwrap();
1483 assert_eq!(response.status(), 200);
1484 mock.assert_async().await;
1485 }
1486
1487 #[tokio::test]
1489 async fn new_databricks_applies_string_options() {
1490 let mut server = mockito::Server::new_async().await;
1491 let mock = server
1492 .mock("GET", "/ping")
1493 .match_header("authorization", "Bearer option-pat")
1494 .with_status(200)
1495 .create_async()
1496 .await;
1497
1498 let client = CloudClient::new_databricks(
1499 [
1500 ("databricks_host", server.url()),
1501 ("databricks_token", "option-pat".to_string()),
1502 ("allow_http", "true".to_string()),
1503 ("not_a_real_key", "ignored".to_string()),
1504 ],
1505 None,
1506 )
1507 .unwrap();
1508
1509 let response = client
1510 .get(format!("{}/ping", server.url()))
1511 .send()
1512 .await
1513 .unwrap();
1514 assert_eq!(response.status(), 200);
1515 mock.assert_async().await;
1516 }
1517}
1518
1519#[cfg(test)]
1520mod databricks_azure_auth_tests {
1521 use super::*;
1522 use databricks::{DatabricksBuilder, DatabricksConfigKey};
1523
1524 fn azure_sp_builder(server: &mockito::Server) -> DatabricksBuilder {
1525 DatabricksBuilder::new()
1526 .with_config(DatabricksConfigKey::Host, server.url())
1527 .with_config(DatabricksConfigKey::ClientId, "sp-id")
1528 .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1529 .with_config(DatabricksConfigKey::AzureTenantId, "tenant-1")
1530 .with_config(DatabricksConfigKey::AzureAuthorityHost, server.url())
1531 .with_config(
1532 DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1533 "true",
1534 )
1535 }
1536
1537 #[tokio::test]
1540 async fn tenant_selects_azure_client_secret_auth() {
1541 let mut server = mockito::Server::new_async().await;
1542 let aad = server
1543 .mock("POST", "/tenant-1/oauth2/v2.0/token")
1544 .match_body(mockito::Matcher::AllOf(vec![
1545 mockito::Matcher::UrlEncoded("client_id".into(), "sp-id".into()),
1546 mockito::Matcher::UrlEncoded(
1547 "scope".into(),
1548 "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default".into(),
1549 ),
1550 ]))
1551 .with_status(200)
1552 .with_header("content-type", "application/json")
1553 .with_body(r#"{"access_token":"aad-token","expires_in":3600,"token_type":"Bearer"}"#)
1554 .expect(1)
1555 .create_async()
1556 .await;
1557 let m2m = server
1558 .mock("POST", "/oidc/v1/token")
1559 .expect(0)
1560 .create_async()
1561 .await;
1562 let api = server
1563 .mock("GET", "/api/2.1/unity-catalog/catalogs")
1564 .match_header("authorization", "Bearer aad-token")
1565 .with_status(200)
1566 .create_async()
1567 .await;
1568
1569 let client =
1570 CloudClient::new_databricks_from_builder(azure_sp_builder(&server), None).unwrap();
1571 let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1572 assert_eq!(client.get(&url).send().await.unwrap().status(), 200);
1573
1574 aad.assert_async().await;
1575 m2m.assert_async().await;
1576 api.assert_async().await;
1577 }
1578
1579 #[test]
1580 fn forced_azure_client_secret_requires_a_tenant() {
1581 let result = DatabricksBuilder::new()
1582 .with_config(DatabricksConfigKey::ClientId, "sp-id")
1583 .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1584 .with_auth_type("azure-client-secret")
1585 .build(None);
1586 let Err(err) = result else {
1587 panic!("expected an error without a tenant id")
1588 };
1589 assert!(err.to_string().contains("AZURE_TENANT_ID"), "{err}");
1590 }
1591
1592 #[test]
1593 fn azure_config_keys_roundtrip() {
1594 for key in [
1595 DatabricksConfigKey::AzureTenantId,
1596 DatabricksConfigKey::AzureAuthorityHost,
1597 ] {
1598 assert_eq!(key.as_ref().parse::<DatabricksConfigKey>().unwrap(), key);
1599 }
1600 assert_eq!(
1601 "azure_tenant_id".parse::<DatabricksConfigKey>().unwrap(),
1602 DatabricksConfigKey::AzureTenantId
1603 );
1604 }
1605}