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();
872 let temp_path = temp_dir.path().to_path_buf();
873
874 let mut server = mockito::Server::new_async().await;
875 let mock = server
876 .mock("GET", "/test")
877 .with_status(200)
878 .with_header("content-type", "application/json")
879 .with_body(r#"{"message": "Hello, World!"}"#)
880 .create_async()
881 .await;
882
883 let mut client = CloudClient::new_unauthenticated();
884 client.set_recording_dir(temp_path.clone()).unwrap();
885
886 let url = format!("{}/test", server.url());
887 let response = client.get(&url).send().await.unwrap();
888
889 assert_eq!(response.status(), 200);
890 let body = response.text().await.unwrap();
891 assert_eq!(body, r#"{"message": "Hello, World!"}"#);
892
893 let recordings: Vec<_> = fs::read_dir(&temp_path)
894 .unwrap()
895 .filter_map(|entry| {
896 let entry = entry.ok()?;
897 let path = entry.path();
898 if path.extension()? == "json" {
899 Some(path)
900 } else {
901 None
902 }
903 })
904 .collect();
905
906 assert_eq!(recordings.len(), 1, "Expected exactly one recording file");
907
908 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
909 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
910
911 assert_eq!(recording.request.method, "GET");
912 assert_eq!(recording.request.url_path, "/test");
913 assert_eq!(recording.request.body, None);
914
915 assert_eq!(recording.response.status, 200);
916 assert_eq!(
917 recording.response.headers.get("content-type").unwrap(),
918 "application/json"
919 );
920 assert_eq!(
921 recording.response.body.as_ref().unwrap(),
922 r#"{"message": "Hello, World!"}"#
923 );
924
925 mock.assert_async().await;
926 }
927
928 #[tokio::test]
929 async fn test_recording_with_request_body() {
930 let temp_dir = TempDir::new().unwrap();
931 let temp_path = temp_dir.path().to_path_buf();
932
933 let mut server = mockito::Server::new_async().await;
934 let mock = server
935 .mock("POST", "/create")
936 .with_status(201)
937 .with_header("location", "/resource/123")
938 .with_body(r#"{"id": 123, "status": "created"}"#)
939 .create_async()
940 .await;
941
942 let mut client = CloudClient::new_unauthenticated();
943 client.set_recording_dir(temp_path.clone()).unwrap();
944
945 let url = format!("{}/create", server.url());
946 let response = client
947 .post(&url)
948 .json(&serde_json::json!({"name": "test resource"}))
949 .send()
950 .await
951 .unwrap();
952
953 assert_eq!(response.status(), 201);
954
955 let recordings: Vec<_> = fs::read_dir(&temp_path)
956 .unwrap()
957 .filter_map(|entry| {
958 let entry = entry.ok()?;
959 let path = entry.path();
960 if path.extension()? == "json" {
961 Some(path)
962 } else {
963 None
964 }
965 })
966 .collect();
967
968 assert_eq!(recordings.len(), 1);
969
970 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
971 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
972
973 assert_eq!(recording.request.method, "POST");
974 assert_eq!(recording.request.url_path, "/create");
975 assert!(recording.request.body.is_some());
976 assert!(recording.request.body.unwrap().contains("test resource"));
977
978 assert_eq!(recording.response.status, 201);
979 assert_eq!(
980 recording.response.headers.get("location").unwrap(),
981 "/resource/123"
982 );
983 assert!(recording.response.body.unwrap().contains("created"));
984
985 mock.assert_async().await;
986 }
987
988 #[tokio::test]
989 async fn test_counter_based_file_naming() {
990 let temp_dir = TempDir::new().unwrap();
991 let temp_path = temp_dir.path().to_path_buf();
992
993 let mut server = mockito::Server::new_async().await;
994 let mock1 = server
995 .mock("GET", "/first")
996 .with_status(200)
997 .with_body("first response")
998 .create_async()
999 .await;
1000 let mock2 = server
1001 .mock("GET", "/second")
1002 .with_status(200)
1003 .with_body("second response")
1004 .create_async()
1005 .await;
1006
1007 let mut client = CloudClient::new_unauthenticated();
1008 client.set_recording_dir(temp_path.clone()).unwrap();
1009
1010 let url1 = format!("{}/first", server.url());
1011 let url2 = format!("{}/second", server.url());
1012
1013 let _response1 = client.get(&url1).send().await.unwrap();
1014 let _response2 = client.get(&url2).send().await.unwrap();
1015
1016 let mut recordings: Vec<_> = fs::read_dir(&temp_path)
1017 .unwrap()
1018 .filter_map(|entry| {
1019 let entry = entry.ok()?;
1020 let path = entry.path();
1021 if path.extension()? == "json" {
1022 Some(path)
1023 } else {
1024 None
1025 }
1026 })
1027 .collect();
1028
1029 recordings.sort();
1030 assert_eq!(recordings.len(), 2);
1031
1032 assert!(recordings[0].file_name().unwrap().to_str().unwrap() == "0000.json");
1033 assert!(recordings[1].file_name().unwrap().to_str().unwrap() == "0001.json");
1034
1035 let first_content = fs::read_to_string(&recordings[0]).unwrap();
1036 let first_recording: RequestResponseInfo = serde_json::from_str(&first_content).unwrap();
1037 assert_eq!(first_recording.request.url_path, "/first");
1038 assert_eq!(
1039 first_recording.response.body.as_ref().unwrap(),
1040 "first response"
1041 );
1042
1043 let second_content = fs::read_to_string(&recordings[1]).unwrap();
1044 let second_recording: RequestResponseInfo = serde_json::from_str(&second_content).unwrap();
1045 assert_eq!(second_recording.request.url_path, "/second");
1046 assert_eq!(
1047 second_recording.response.body.as_ref().unwrap(),
1048 "second response"
1049 );
1050
1051 mock1.assert_async().await;
1052 mock2.assert_async().await;
1053 }
1054
1055 #[tokio::test]
1056 async fn test_query_parameter_recording() {
1057 let temp_dir = TempDir::new().unwrap();
1058 let temp_path = temp_dir.path().to_path_buf();
1059
1060 let mut server = mockito::Server::new_async().await;
1061
1062 let mock = server
1063 .mock("GET", "/catalogs?max_results=10&page_token=abc123")
1064 .with_status(200)
1065 .with_header("content-type", "application/json")
1066 .with_body(r#"{"catalogs": []}"#)
1067 .create_async()
1068 .await;
1069
1070 let mut client = CloudClient::new_unauthenticated();
1071 client.set_recording_dir(temp_path.clone()).unwrap();
1072
1073 let url = format!("{}/catalogs?max_results=10&page_token=abc123", server.url());
1074 let response = client.get(&url).send().await.unwrap();
1075
1076 assert!(response.status().is_success());
1077
1078 let recordings: Vec<_> = fs::read_dir(&temp_path)
1079 .unwrap()
1080 .filter_map(|entry| {
1081 let entry = entry.ok()?;
1082 let path = entry.path();
1083 if path.extension()? == "json" {
1084 Some(path)
1085 } else {
1086 None
1087 }
1088 })
1089 .collect();
1090
1091 assert_eq!(recordings.len(), 1);
1092
1093 let recording_content = fs::read_to_string(&recordings[0]).unwrap();
1095 let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
1096
1097 assert_eq!(recording.request.method, "GET");
1099 assert_eq!(
1100 recording.request.url_path,
1101 "/catalogs?max_results=10&page_token=abc123"
1102 );
1103 assert_eq!(recording.request.body, None);
1104
1105 assert_eq!(recording.response.status, 200);
1107 assert_eq!(
1108 recording.response.body.as_ref().unwrap(),
1109 r#"{"catalogs": []}"#
1110 );
1111
1112 mock.assert_async().await;
1113 }
1114
1115 #[tokio::test]
1120 async fn test_recording_does_not_contain_bearer_token_value() {
1121 let temp_dir = TempDir::new().unwrap();
1122 let temp_path = temp_dir.path().to_path_buf();
1123
1124 let mut server = mockito::Server::new_async().await;
1125 let mock = server
1126 .mock("GET", "/secret")
1127 .match_header("authorization", mockito::Matcher::Any)
1128 .with_status(200)
1129 .with_body("ok")
1131 .create_async()
1132 .await;
1133
1134 let mut client = CloudClient::new_with_token("super-secret-token-12345");
1135 client.set_recording_dir(temp_path.clone()).unwrap();
1136
1137 let url = format!("{}/secret", server.url());
1138 client.get(&url).send().await.unwrap();
1139
1140 let recording_path = temp_path.join("0000.json");
1141 let content = fs::read_to_string(&recording_path).unwrap();
1142
1143 assert!(
1145 !content.contains("super-secret-token-12345"),
1146 "raw bearer token leaked into recording: {content}"
1147 );
1148
1149 mock.assert_async().await;
1150 }
1151
1152 #[tokio::test]
1156 async fn test_recording_redacts_sensitive_response_headers() {
1157 let temp_dir = TempDir::new().unwrap();
1158 let temp_path = temp_dir.path().to_path_buf();
1159
1160 let mut server = mockito::Server::new_async().await;
1161 let mock = server
1163 .mock("GET", "/s3")
1164 .with_status(200)
1165 .with_header("x-amz-security-token", "AQoXnyc4LLI2AJvUAMOGAR8a1234567890")
1166 .with_header("authorization", "Bearer should-be-redacted")
1167 .with_body("{}")
1168 .create_async()
1169 .await;
1170
1171 let mut client = CloudClient::new_unauthenticated();
1172 client.set_recording_dir(temp_path.clone()).unwrap();
1173
1174 let url = format!("{}/s3", server.url());
1175 client.get(&url).send().await.unwrap();
1176
1177 let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1178 assert!(
1179 !content.contains("AQoXnyc4LLI2AJvUAMOGAR8a1234567890"),
1180 "x-amz-security-token leaked into recording: {content}"
1181 );
1182 assert!(
1183 !content.contains("should-be-redacted"),
1184 "Authorization value leaked into recording: {content}"
1185 );
1186 assert_eq!(
1188 content.matches("<REDACTED>").count(),
1189 2,
1190 "expected 2 <REDACTED> entries in recording: {content}"
1191 );
1192
1193 mock.assert_async().await;
1194 }
1195
1196 #[tokio::test]
1198 async fn test_recording_remains_valid_json_after_redaction() {
1199 let temp_dir = TempDir::new().unwrap();
1200 let temp_path = temp_dir.path().to_path_buf();
1201
1202 let mut server = mockito::Server::new_async().await;
1203 let _mock = server
1204 .mock("GET", "/check")
1205 .with_status(200)
1206 .with_header("authorization", "Bearer should-be-redacted")
1207 .with_body(r#"{"ok":true}"#)
1208 .create_async()
1209 .await;
1210
1211 let mut client = CloudClient::new_unauthenticated();
1212 client.set_recording_dir(temp_path.clone()).unwrap();
1213
1214 let url = format!("{}/check", server.url());
1215 client.get(&url).send().await.unwrap();
1216
1217 let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1218 let parsed: RequestResponseInfo = serde_json::from_str(&content)
1220 .expect("recording file must be valid JSON even after redaction");
1221 assert_eq!(parsed.response.status, 200);
1222 }
1223}
1224
1225#[cfg(all(test, not(feature = "recording")))]
1230mod retry_integration_tests {
1231 use super::*;
1232 use std::time::Duration;
1233
1234 fn fast_retry(max_retries: usize) -> RetryConfig {
1235 RetryConfig {
1236 backoff: crate::backoff::BackoffConfig {
1237 init_backoff: Duration::from_millis(1),
1238 max_backoff: Duration::from_millis(5),
1239 base: 2.,
1240 },
1241 max_retries,
1242 retry_timeout: Duration::from_secs(30),
1243 }
1244 }
1245
1246 #[tokio::test]
1249 async fn send_retries_server_error_then_succeeds() {
1250 let mut server = mockito::Server::new_async().await;
1251 let fail = server
1252 .mock("GET", "/r")
1253 .with_status(503)
1254 .expect(1)
1255 .create_async()
1256 .await;
1257 let ok = server
1258 .mock("GET", "/r")
1259 .with_status(200)
1260 .with_body("ok")
1261 .expect(1)
1262 .create_async()
1263 .await;
1264
1265 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1266 let resp = client
1267 .get(format!("{}/r", server.url()))
1268 .send()
1269 .await
1270 .unwrap();
1271
1272 assert_eq!(resp.status(), 200);
1273 assert_eq!(resp.text().await.unwrap(), "ok");
1274 fail.assert_async().await;
1275 ok.assert_async().await;
1276 }
1277
1278 #[tokio::test]
1281 async fn send_does_not_retry_client_error() {
1282 let mut server = mockito::Server::new_async().await;
1283 let mock = server
1284 .mock("GET", "/r")
1285 .with_status(404)
1286 .expect(1)
1287 .create_async()
1288 .await;
1289
1290 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1291 let err = client
1292 .get(format!("{}/r", server.url()))
1293 .send()
1294 .await
1295 .unwrap_err();
1296
1297 assert!(matches!(err, Error::NotFound { .. }), "got {err:?}");
1298 mock.assert_async().await;
1299 }
1300
1301 #[tokio::test]
1306 async fn send_raw_preserves_status_and_body() {
1307 let mut server = mockito::Server::new_async().await;
1308
1309 let not_found = server
1311 .mock("GET", "/nf")
1312 .with_status(404)
1313 .with_body(r#"{"error":{"type":"NoSuchTableException"}}"#)
1314 .expect(1)
1315 .create_async()
1316 .await;
1317 let conflict = server
1319 .mock("POST", "/cf")
1320 .with_status(409)
1321 .with_body(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1322 .expect(1)
1323 .create_async()
1324 .await;
1325 let fail = server
1327 .mock("GET", "/rt")
1328 .with_status(503)
1329 .expect(1)
1330 .create_async()
1331 .await;
1332 let ok = server
1333 .mock("GET", "/rt")
1334 .with_status(200)
1335 .with_body("ok")
1336 .expect(1)
1337 .create_async()
1338 .await;
1339
1340 let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1341
1342 let err = client
1343 .get(format!("{}/nf", server.url()))
1344 .send_raw()
1345 .await
1346 .unwrap_err();
1347 match err {
1348 SendRawError::Retry(e) => {
1349 assert_eq!(e.status(), Some(reqwest::StatusCode::NOT_FOUND));
1350 assert_eq!(
1351 e.body(),
1352 Some(r#"{"error":{"type":"NoSuchTableException"}}"#)
1353 );
1354 }
1355 other => panic!("expected Retry error, got {other:?}"),
1356 }
1357
1358 let err = client
1359 .post(format!("{}/cf", server.url()))
1360 .send_raw()
1361 .await
1362 .unwrap_err();
1363 match err {
1364 SendRawError::Retry(e) => {
1365 assert_eq!(e.status(), Some(reqwest::StatusCode::CONFLICT));
1366 assert_eq!(
1367 e.body(),
1368 Some(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1369 );
1370 }
1371 other => panic!("expected Retry error, got {other:?}"),
1372 }
1373
1374 let resp = client
1375 .get(format!("{}/rt", server.url()))
1376 .send_raw()
1377 .await
1378 .unwrap();
1379 assert_eq!(resp.status(), 200);
1380 assert_eq!(resp.text().await.unwrap(), "ok");
1381
1382 not_found.assert_async().await;
1383 conflict.assert_async().await;
1384 fail.assert_async().await;
1385 ok.assert_async().await;
1386 }
1387
1388 #[tokio::test]
1391 async fn sign_and_send_signs_and_retries() {
1392 let mut server = mockito::Server::new_async().await;
1393 let fail = server
1394 .mock("POST", "/rpc")
1395 .match_header("authorization", "Bearer tok")
1396 .with_status(503)
1397 .expect(1)
1398 .create_async()
1399 .await;
1400 let ok = server
1401 .mock("POST", "/rpc")
1402 .match_header("authorization", "Bearer tok")
1403 .with_status(200)
1404 .with_body("pong")
1405 .expect(1)
1406 .create_async()
1407 .await;
1408
1409 let client = CloudClient::new_with_token("tok").with_retry_config(fast_retry(3));
1410 let request = reqwest::Client::new()
1411 .post(format!("{}/rpc", server.url()))
1412 .body("ping")
1413 .build()
1414 .unwrap();
1415 let resp = client.sign_and_send(request).await.unwrap();
1416
1417 assert_eq!(resp.status(), 200);
1418 assert_eq!(resp.text().await.unwrap(), "pong");
1419 fail.assert_async().await;
1420 ok.assert_async().await;
1421 }
1422}
1423
1424#[cfg(test)]
1425mod databricks_client_tests {
1426 use super::*;
1427 use databricks::{DatabricksBuilder, DatabricksConfigKey};
1428
1429 #[tokio::test]
1431 async fn new_databricks_from_builder_signs_with_builder_credential() {
1432 let mut server = mockito::Server::new_async().await;
1433 let mock = server
1434 .mock("GET", "/api/2.1/unity-catalog/catalogs")
1435 .match_header("authorization", "Bearer builder-pat")
1436 .with_status(200)
1437 .create_async()
1438 .await;
1439
1440 let builder = DatabricksBuilder::new()
1441 .with_config(DatabricksConfigKey::Host, server.url())
1442 .with_config(DatabricksConfigKey::Token, "builder-pat")
1443 .with_config(
1444 DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1445 "true",
1446 );
1447 let client = CloudClient::new_databricks_from_builder(builder, None).unwrap();
1448
1449 let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1450 let response = client.get(&url).send().await.unwrap();
1451 assert_eq!(response.status(), 200);
1452 mock.assert_async().await;
1453 }
1454
1455 #[tokio::test]
1457 async fn new_databricks_applies_string_options() {
1458 let mut server = mockito::Server::new_async().await;
1459 let mock = server
1460 .mock("GET", "/ping")
1461 .match_header("authorization", "Bearer option-pat")
1462 .with_status(200)
1463 .create_async()
1464 .await;
1465
1466 let client = CloudClient::new_databricks(
1467 [
1468 ("databricks_host", server.url()),
1469 ("databricks_token", "option-pat".to_string()),
1470 ("allow_http", "true".to_string()),
1471 ("not_a_real_key", "ignored".to_string()),
1472 ],
1473 None,
1474 )
1475 .unwrap();
1476
1477 let response = client
1478 .get(format!("{}/ping", server.url()))
1479 .send()
1480 .await
1481 .unwrap();
1482 assert_eq!(response.status(), 200);
1483 mock.assert_async().await;
1484 }
1485}
1486
1487#[cfg(test)]
1488mod databricks_azure_auth_tests {
1489 use super::*;
1490 use databricks::{DatabricksBuilder, DatabricksConfigKey};
1491
1492 fn azure_sp_builder(server: &mockito::Server) -> DatabricksBuilder {
1493 DatabricksBuilder::new()
1494 .with_config(DatabricksConfigKey::Host, server.url())
1495 .with_config(DatabricksConfigKey::ClientId, "sp-id")
1496 .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1497 .with_config(DatabricksConfigKey::AzureTenantId, "tenant-1")
1498 .with_config(DatabricksConfigKey::AzureAuthorityHost, server.url())
1499 .with_config(
1500 DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1501 "true",
1502 )
1503 }
1504
1505 #[tokio::test]
1508 async fn tenant_selects_azure_client_secret_auth() {
1509 let mut server = mockito::Server::new_async().await;
1510 let aad = server
1511 .mock("POST", "/tenant-1/oauth2/v2.0/token")
1512 .match_body(mockito::Matcher::AllOf(vec![
1513 mockito::Matcher::UrlEncoded("client_id".into(), "sp-id".into()),
1514 mockito::Matcher::UrlEncoded(
1515 "scope".into(),
1516 "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default".into(),
1517 ),
1518 ]))
1519 .with_status(200)
1520 .with_header("content-type", "application/json")
1521 .with_body(r#"{"access_token":"aad-token","expires_in":3600,"token_type":"Bearer"}"#)
1522 .expect(1)
1523 .create_async()
1524 .await;
1525 let m2m = server
1526 .mock("POST", "/oidc/v1/token")
1527 .expect(0)
1528 .create_async()
1529 .await;
1530 let api = server
1531 .mock("GET", "/api/2.1/unity-catalog/catalogs")
1532 .match_header("authorization", "Bearer aad-token")
1533 .with_status(200)
1534 .create_async()
1535 .await;
1536
1537 let client =
1538 CloudClient::new_databricks_from_builder(azure_sp_builder(&server), None).unwrap();
1539 let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1540 assert_eq!(client.get(&url).send().await.unwrap().status(), 200);
1541
1542 aad.assert_async().await;
1543 m2m.assert_async().await;
1544 api.assert_async().await;
1545 }
1546
1547 #[test]
1548 fn forced_azure_client_secret_requires_a_tenant() {
1549 let result = DatabricksBuilder::new()
1550 .with_config(DatabricksConfigKey::ClientId, "sp-id")
1551 .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1552 .with_auth_type("azure-client-secret")
1553 .build(None);
1554 let Err(err) = result else {
1555 panic!("expected an error without a tenant id")
1556 };
1557 assert!(err.to_string().contains("AZURE_TENANT_ID"), "{err}");
1558 }
1559
1560 #[test]
1561 fn azure_config_keys_roundtrip() {
1562 for key in [
1563 DatabricksConfigKey::AzureTenantId,
1564 DatabricksConfigKey::AzureAuthorityHost,
1565 ] {
1566 assert_eq!(key.as_ref().parse::<DatabricksConfigKey>().unwrap(), key);
1567 }
1568 assert_eq!(
1569 "azure_tenant_id".parse::<DatabricksConfigKey>().unwrap(),
1570 DatabricksConfigKey::AzureTenantId
1571 );
1572 }
1573}