Skip to main content

olai_http/
lib.rs

1//! Unified cloud credential abstraction and HTTP client for AWS, Azure, GCP, and
2//! Databricks.
3//!
4//! This crate provides a single authenticated HTTP client, [`CloudClient`], that
5//! signs outgoing requests for any supported cloud provider. Rather than pulling
6//! in a separate vendor SDK for each cloud, a service constructs one
7//! [`CloudClient`] per provider and issues requests through a familiar,
8//! `reqwest`-style builder ([`CloudRequestBuilder`]). Every provider is reached
9//! through the same [`RequestSigner`] trait, so credential resolution, token
10//! refresh, and request signing are uniform across clouds. The credential
11//! machinery is extracted from the
12//! [`object_store`](https://crates.io/crates/object_store) crate's internal client.
13//!
14//! # Providers
15//!
16//! Each provider has a builder under its own module and a matching
17//! [`CloudClient`] constructor:
18//!
19//! - **AWS** ([`aws`]) — SigV4 signing with static keys, IMDS, ECS/EKS task
20//!   roles, web identity, and STS `AssumeRole`. See [`CloudClient::new_aws`].
21//! - **Azure** ([`azure`]) — Azure AD bearer tokens via client secret, managed
22//!   identity, workload identity, or the Azure CLI. See [`CloudClient::new_azure`].
23//! - **Google Cloud** ([`gcp`]) — OAuth 2.0 bearer tokens via service-account
24//!   JWTs, the GCE metadata server, or workload identity federation. See
25//!   [`CloudClient::new_google`].
26//! - **Databricks** ([`databricks`]) — OAuth M2M and OIDC token exchange. See
27//!   [`CloudClient::new_databricks`].
28//!
29//! For a static token or no authentication at all, use
30//! [`CloudClient::new_with_token`] or [`CloudClient::new_unauthenticated`].
31//!
32//! # Examples
33//!
34//! ```no_run
35//! use olai_http::CloudClient;
36//!
37//! # async fn run() -> olai_http::Result<()> {
38//! let client = CloudClient::new_with_token("my-token");
39//! let resp = client
40//!     .get("https://api.example.com/data")
41//!     .send()
42//!     .await?;
43//! println!("status: {}", resp.status());
44//! # Ok(())
45//! # }
46//! ```
47//!
48//! Enable the `recording` feature to capture HTTP interactions to JSON (with
49//! sensitive headers redacted) for test replay.
50
51#[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
94/// A shared, cloneable signer for a `CloudClient`.
95type SharedSigner = Arc<dyn RequestSigner>;
96
97/// A no-op signer used by `new_unauthenticated`.
98#[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/// A signer that injects a static bearer token.
111#[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/// An authenticated HTTP client for cloud provider APIs.
134///
135/// Created via the provider-specific constructors [`CloudClient::new_aws`],
136/// [`CloudClient::new_azure`], [`CloudClient::new_google`], or
137/// [`CloudClient::new_databricks`], or the simpler [`CloudClient::new_with_token`]
138/// and [`CloudClient::new_unauthenticated`].
139#[derive(Clone)]
140pub struct CloudClient {
141    signer: SharedSigner,
142    reqwest_client: Client,
143    service: Arc<dyn HttpService>,
144    /// Retry configuration applied to requests sent through this client.
145    ///
146    /// Used both by credential providers (token refresh) and by user-initiated
147    /// requests via [`CloudRequestBuilder::send`] and
148    /// [`CloudClient::sign_and_send`]. Override with
149    /// [`CloudClient::with_retry_config`].
150    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    /// Create a new client with AWS credentials.
173    ///
174    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
175    /// will be spawned on the given runtime handle.
176    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    /// Create a new client with Google Cloud credentials.
205    ///
206    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
207    /// will be spawned on the given runtime handle.
208    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    /// Create a new client with Azure credentials.
237    ///
238    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
239    /// will be spawned on the given runtime handle.
240    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    /// Create a new client with Databricks credentials.
269    ///
270    /// Options are applied to a fresh [`databricks::DatabricksBuilder`]; unknown
271    /// keys are ignored. Use [`new_databricks_from_builder`] to start from
272    /// [`databricks::DatabricksBuilder::from_env`] instead, which also loads
273    /// `.databrickscfg` profiles.
274    ///
275    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
276    /// will be spawned on the given runtime handle.
277    ///
278    /// [`new_databricks_from_builder`]: CloudClient::new_databricks_from_builder
279    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    /// Create a new client with Databricks credentials from a configured builder.
296    ///
297    /// Unlike [`new_databricks`], this keeps everything the builder already
298    /// resolved — notably the `.databrickscfg` profile values loaded by
299    /// [`databricks::DatabricksBuilder::from_env`].
300    ///
301    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
302    /// will be spawned on the given runtime handle.
303    ///
304    /// [`new_databricks`]: CloudClient::new_databricks
305    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    /// Create a client that signs requests with tokens from `provider`.
323    ///
324    /// See [`bearer::BearerAuth`] for providers that work against any API, not
325    /// just Databricks. `options` and `retry` configure requests sent through
326    /// this client; if `runtime` is provided, their I/O is spawned on it.
327    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    /// Create a new client with a personal access token.
344    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    /// Create a new unauthenticated client.
358    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    /// Create an unauthenticated client with explicit HTTP and retry settings.
370    ///
371    /// Like [`new_unauthenticated`](Self::new_unauthenticated), for anonymous calls that
372    /// need timeouts, proxies or a bounded retry budget, such as
373    /// [`aws::bucket_region`]. If `runtime` is provided, all HTTP I/O is spawned on it.
374    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    /// Route all HTTP I/O through the given runtime handle.
390    ///
391    /// This is useful for simple constructors (`new_with_token`, `new_unauthenticated`)
392    /// where no credential providers need to perform HTTP I/O. For cloud provider
393    /// constructors, pass the handle at construction time instead.
394    pub fn with_runtime(mut self, handle: Handle) -> Self {
395        self.service = Arc::new(SpawnService::new(self.service, handle));
396        self
397    }
398
399    /// Replace the [`HttpService`] used for request execution.
400    pub fn with_http_service(mut self, service: Arc<dyn HttpService>) -> Self {
401        self.service = service;
402        self
403    }
404
405    /// Override the [`RetryConfig`] applied to requests sent through this client.
406    ///
407    /// Requests issued via [`CloudRequestBuilder::send`] and
408    /// [`sign_and_send`](Self::sign_and_send) are retried per this config
409    /// (exponential backoff with jitter; safe/idempotent requests are also
410    /// retried on timeout). The [default](RetryConfig::default) allows up to 10
411    /// retries — lower it for latency-sensitive interactive calls.
412    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
413        self.retry_config = retry_config;
414        self
415    }
416
417    /// Sign an already-built [`reqwest::Request`] and dispatch it, with retries.
418    ///
419    /// This is the request-level counterpart to the [`CloudRequestBuilder`]
420    /// flow: it applies the client's [`RequestSigner`] (refreshing credentials
421    /// as needed) and sends the request through the configured [`HttpService`],
422    /// retrying transient failures per the client's [`RetryConfig`]. It is
423    /// useful when the request is produced elsewhere (for example by a protocol
424    /// stack such as ConnectRPC) and only needs authentication and transport.
425    ///
426    /// The request body must be in memory (not a stream): retries clone the
427    /// request, and provider signers such as AWS SigV4 hash the body to sign it.
428    ///
429    /// # Errors
430    ///
431    /// Returns an [`Error`] if signing fails or if the request ultimately fails
432    /// after exhausting retries. A non-success HTTP status code is not itself an
433    /// error; inspect [`reqwest::Response::status`] on the returned response.
434    pub async fn sign_and_send(&self, request: reqwest::Request) -> Result<reqwest::Response> {
435        // The signer operates on a `RequestBuilder`; reconstruct one from the
436        // pre-built request (cloning the in-memory body) so the existing
437        // sign + retry path applies unchanged.
438        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    /// Start building a request for the given HTTP `method` and `url`.
454    ///
455    /// The returned [`CloudRequestBuilder`] borrows this client's signer, so the
456    /// request is signed for the configured provider when
457    /// [`CloudRequestBuilder::send`] is called.
458    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    /// Start building a `GET` request for `url`. Shortcut for [`request`](Self::request).
468    pub fn get<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
469        self.request(Method::GET, url)
470    }
471
472    /// Start building a `POST` request for `url`. Shortcut for [`request`](Self::request).
473    pub fn post<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
474        self.request(Method::POST, url)
475    }
476
477    /// Start building a `PUT` request for `url`. Shortcut for [`request`](Self::request).
478    pub fn put<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
479        self.request(Method::PUT, url)
480    }
481
482    /// Start building a `DELETE` request for `url`. Shortcut for [`request`](Self::request).
483    pub fn delete<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
484        self.request(Method::DELETE, url)
485    }
486
487    /// Start building a `HEAD` request for `url`. Shortcut for [`request`](Self::request).
488    pub fn head<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
489        self.request(Method::HEAD, url)
490    }
491
492    /// Start building a `PATCH` request for `url`. Shortcut for [`request`](Self::request).
493    pub fn patch<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
494        self.request(Method::PATCH, url)
495    }
496
497    /// Start building an `OPTIONS` request for `url`. Shortcut for [`request`](Self::request).
498    pub fn options<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
499        self.request(Method::OPTIONS, url)
500    }
501
502    /// Start building a `TRACE` request for `url`. Shortcut for [`request`](Self::request).
503    pub fn trace<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
504        self.request(Method::TRACE, url)
505    }
506
507    /// Start building a `CONNECT` request for `url`. Shortcut for [`request`](Self::request).
508    pub fn connect<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
509        self.request(Method::CONNECT, url)
510    }
511
512    /// Enable request/response recording, writing each interaction to `out_dir`.
513    ///
514    /// Once set, every request sent through this client is captured to a
515    /// numbered JSON file (`0000.json`, `0001.json`, …) under `out_dir` for
516    /// later test replay. Sensitive response headers (`authorization`,
517    /// `x-amz-security-token`, `cookie`, and similar) are replaced with
518    /// `"<REDACTED>"` before anything is written to disk, so
519    /// recordings never persist bearer tokens or signing secrets. Request
520    /// headers are not recorded at all.
521    ///
522    /// `out_dir` is canonicalized eagerly, so it must already exist.
523    ///
524    /// Only available when the `recording` feature is enabled.
525    ///
526    /// # Errors
527    ///
528    /// Returns the [`io::Error`](std::io::Error) from canonicalizing `out_dir`,
529    /// for example if the directory does not exist or is not accessible.
530    #[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
541/// A builder for a single request issued through a [`CloudClient`].
542///
543/// Created by [`CloudClient::request`] and the per-verb shortcuts such as
544/// [`CloudClient::get`]. Configure the request with the builder methods, then
545/// call [`send`](Self::send) to sign and dispatch it.
546pub struct CloudRequestBuilder {
547    builder: RequestBuilder,
548    client: CloudClient,
549    #[cfg(feature = "recording")]
550    out_dir: Option<PathBuf>,
551}
552
553impl CloudRequestBuilder {
554    /// Add a `Header` to this Request.
555    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    /// Add a set of Headers to the existing ones on this Request.
567    ///
568    /// The headers will be merged in to any already set.
569    pub fn headers(mut self, headers: HeaderMap) -> CloudRequestBuilder {
570        self.builder = self.builder.headers(headers);
571        self
572    }
573
574    /// Set the request body.
575    pub fn body<T: Into<Body>>(mut self, body: T) -> CloudRequestBuilder {
576        self.builder = self.builder.body(body);
577        self
578    }
579
580    /// Enables a request timeout.
581    ///
582    /// The timeout is applied from when the request starts connecting until the
583    /// response body has finished. It affects only this request and overrides
584    /// the timeout configured using `ClientBuilder::timeout()`.
585    pub fn timeout(mut self, timeout: Duration) -> CloudRequestBuilder {
586        self.builder = self.builder.timeout(timeout);
587        self
588    }
589
590    /// Modify the query string of the URL.
591    ///
592    /// Modifies the URL of this request, adding the parameters provided.
593    /// This method appends and does not overwrite. This means that it can
594    /// be called multiple times and that existing query parameters are not
595    /// overwritten if the same key is used. The key will simply show up
596    /// twice in the query string.
597    /// Calling `.query(&[("foo", "a"), ("foo", "b")])` gives `"foo=a&foo=b"`.
598    ///
599    /// # Note
600    /// This method does not support serializing a single key-value
601    /// pair. Instead of using `.query(("key", "val"))`, use a sequence, such
602    /// as `.query(&[("key", "val")])`. It's also possible to serialize structs
603    /// and maps into a key-value pair.
604    ///
605    /// # Errors
606    /// This method will fail if the object you provide cannot be serialized
607    /// into a query string.
608    pub fn query<T: Serialize + ?Sized>(mut self, query: &T) -> CloudRequestBuilder {
609        self.builder = self.builder.query(query);
610        self
611    }
612
613    /// Send a JSON body.
614    ///
615    /// # Errors
616    ///
617    /// Serialization can fail if `T`'s implementation of `Serialize` decides to
618    /// fail, or if `T` contains a map with non-string keys.
619    pub fn json<T: Serialize + ?Sized>(mut self, json: &T) -> CloudRequestBuilder {
620        self.builder = self.builder.json(json);
621        self
622    }
623
624    /// Sign and send the request, returning the [`reqwest::Response`].
625    ///
626    /// The request is first passed to the client's [`RequestSigner`], which
627    /// attaches the provider-specific authentication (e.g. AWS SigV4 headers or
628    /// an `Authorization: Bearer` header), refreshing any cached credential as
629    /// needed. The signed request is then dispatched through the client's
630    /// [`HttpService`], retrying transient failures per the client's
631    /// [`RetryConfig`] and mapping a non-success status to the matching
632    /// [`Error`].
633    ///
634    /// When the `recording` feature is enabled and a recording directory has
635    /// been configured via `CloudClient::set_recording_dir`, the request and
636    /// response are instead captured to disk (with sensitive headers redacted)
637    /// in a single round-trip — this capture path does **not** retry or map
638    /// status codes to errors, so use the default build for production traffic.
639    ///
640    /// # Errors
641    ///
642    /// Returns an [`Error`] if signing fails, if the request cannot be built, or
643    /// if the underlying HTTP transport returns an error. A non-success HTTP
644    /// status code is not itself an error; inspect [`reqwest::Response::status`]
645    /// on the returned response.
646    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    /// Sign and send the request, preserving the classified non-success result.
664    ///
665    /// This behaves exactly like [`send`](Self::send) — same signing, same
666    /// [`RetryConfig`] semantics (5xx and transport failures retried, 4xx
667    /// returned immediately) — but a non-success outcome is surfaced as the
668    /// [`RetryError`] the retry layer already classified, rather than being
669    /// remapped through [`RetryError::error`] into an opaque [`Error`].
670    ///
671    /// Use this when the caller needs the HTTP **status** and **response body**
672    /// of a failed request — for example to parse a protocol-specific error
673    /// envelope. [`RetryError::status`] and [`RetryError::body`] expose them;
674    /// [`send`](Self::send) discards both by boxing the status into a coarse
675    /// [`Error`] variant (`NotFound`, `AlreadyExists`, …).
676    ///
677    /// When the `recording` feature is enabled the request is captured to disk
678    /// via the same single-shot path as [`send`](Self::send) and the recorded
679    /// [`reqwest::Response`] is returned for any status (no status→error
680    /// mapping), so a non-success status is `Ok(response)` in that build — the
681    /// caller inspects [`reqwest::Response::status`] itself.
682    ///
683    /// # Errors
684    ///
685    /// Returns [`Error`] if signing fails. Otherwise, on a non-success status
686    /// (in the default build) returns [`RetryError`] carrying the status and
687    /// response body; a transport failure surfaces as
688    /// [`RetryError::Reqwest`]/[`RetryError::Transport`].
689    #[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    /// See [`send_raw`](Self::send_raw). Under the `recording` feature this maps
699    /// to the single-shot capture path and returns the response for any status.
700    #[cfg(feature = "recording")]
701    pub async fn send_raw(self) -> Result<reqwest::Response, SendRawError> {
702        send_record(self).await
703    }
704}
705
706/// The error returned by [`CloudRequestBuilder::send_raw`].
707///
708/// Distinguishes a failure to *sign* the request (an [`Error`] from the signer,
709/// e.g. a credential-refresh failure) from a failure of the *request itself*
710/// (a [`RetryError`] carrying the HTTP status and response body). This lets a
711/// caller parse a protocol error envelope from [`RetryError::body`] while still
712/// propagating signing failures faithfully.
713#[cfg(not(feature = "recording"))]
714#[derive(Debug, thiserror::Error)]
715pub enum SendRawError {
716    /// The request could not be signed (e.g. credential resolution failed).
717    #[error(transparent)]
718    Sign(#[from] Error),
719    /// The request was sent but failed; carries the classified status + body.
720    #[error(transparent)]
721    Retry(RetryError),
722}
723
724/// Under the `recording` feature `send_raw` returns the recorded response for
725/// any status, so the only error it can produce is a signing/capture [`Error`].
726#[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/// Header names whose values must never appear in recording files.
753///
754/// These headers carry bearer tokens, signing secrets, or session tokens.
755/// They are replaced with `"<REDACTED>"` before any recording is written to disk.
756#[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    // Record the response
809    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    // Get response body while preserving it for the caller
817    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            // Redact sensitive headers before writing to disk
829            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    // Return a new response built from the recorded data, using the raw (unredacted) headers
852    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        // Read and verify the recording content includes query parameters
1094        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
1095        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
1096
1097        // Verify request information includes query parameters
1098        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        // Verify response information
1106        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    /// Verify that the bearer token injected by a signed request does not appear
1116    /// in the recording file. Request headers are currently not recorded, so this
1117    /// test confirms the token does not leak via any other path (e.g. echoed back
1118    /// in a response header or response body).
1119    #[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            // Server does NOT echo the token back — simulates a well-behaved API
1130            .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        // The raw token must not appear in the file at all
1144        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    /// Verify that sensitive headers returned by the server (e.g. a reflected
1153    /// Authorization or AWS security token) are redacted before being written
1154    /// to the recording file.
1155    #[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        // Simulate a server that echoes a sensitive header back in its response
1162        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        // Both headers should be replaced with the redaction sentinel
1187        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    /// Verify that a recording produced with redacted headers can still be parsed.
1197    #[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        // Must parse cleanly as a RequestResponseInfo
1219        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// These assert the production send path's retry + error-mapping contract.
1226// The `recording` feature replaces that path with a single-shot capture (no
1227// retry, no status->error mapping — see `send_record`), so they are scoped to
1228// the non-recording build where the behavior under test actually applies.
1229#[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    // A 5xx is retried by `CloudClient::send`: a single 503 followed by a 200
1247    // surfaces the 200.
1248    #[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    // A 4xx is not retryable: it must surface immediately as an error without
1279    // consuming retries (a second mock would go unmatched).
1280    #[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    // send_raw preserves the classified status + body of a non-success
1302    // response instead of remapping it into an opaque Error. A 404 and a 409
1303    // must both surface as Err(RetryError) whose status()/body() are intact,
1304    // and a 5xx must still be retried (503 then 200 -> Ok(200)).
1305    #[tokio::test]
1306    async fn send_raw_preserves_status_and_body() {
1307        let mut server = mockito::Server::new_async().await;
1308
1309        // 404 with a body envelope: status + body must survive.
1310        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        // 409 conflict: same — non-retryable, body preserved.
1318        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        // 5xx is still retried under send_raw: one 503 then a 200.
1326        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    // sign_and_send takes a pre-built reqwest::Request, applies the signer
1389    // (here a bearer token), and retries transient failures just like send.
1390    #[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    /// A pre-configured builder is used as-is: its credential signs requests.
1430    #[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    /// The option-list constructor still parses keys (and ignores unknown ones).
1456    #[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    /// A tenant id selects Azure AD client credentials for the Azure
1506    /// Databricks application instead of Databricks OAuth M2M.
1507    #[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}