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;
72mod client;
73mod config;
74#[cfg(feature = "connectrpc")]
75pub mod connectrpc;
76mod credential;
77pub mod databricks;
78mod error;
79pub mod gcp;
80pub mod service;
81
82mod retry;
83mod token;
84mod util;
85
86pub use client::{Certificate, ClientConfigKey, ClientOptions};
87pub use credential::*;
88pub use error::*;
89pub use retry::Error as RetryError;
90pub use retry::RetryConfig;
91pub use service::{ReqwestService, SpawnService};
92
93/// A shared, cloneable signer for a `CloudClient`.
94type SharedSigner = Arc<dyn RequestSigner>;
95
96/// A no-op signer used by `new_unauthenticated`.
97#[derive(Debug)]
98struct NoopSigner;
99
100impl RequestSigner for NoopSigner {
101    fn sign<'a>(
102        &'a self,
103        req: RequestBuilder,
104    ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
105        Box::pin(async move { Ok(req) })
106    }
107}
108
109/// A signer that injects a static bearer token.
110#[derive(Debug)]
111struct BearerTokenSigner {
112    token: String,
113}
114
115impl RequestSigner for BearerTokenSigner {
116    fn sign<'a>(
117        &'a self,
118        req: RequestBuilder,
119    ) -> futures::future::BoxFuture<'a, Result<RequestBuilder>> {
120        let token = self.token.clone();
121        Box::pin(async move { Ok(req.bearer_auth(&token)) })
122    }
123}
124
125#[cfg(feature = "recording")]
126#[derive(Debug, Clone)]
127struct RecordingState {
128    out_dir: PathBuf,
129    counter: Arc<AtomicU64>,
130}
131
132/// An authenticated HTTP client for cloud provider APIs.
133///
134/// Created via the provider-specific constructors [`CloudClient::new_aws`],
135/// [`CloudClient::new_azure`], [`CloudClient::new_google`], or
136/// [`CloudClient::new_databricks`], or the simpler [`CloudClient::new_with_token`]
137/// and [`CloudClient::new_unauthenticated`].
138#[derive(Clone)]
139pub struct CloudClient {
140    signer: SharedSigner,
141    reqwest_client: Client,
142    service: Arc<dyn HttpService>,
143    /// Retry configuration applied to requests sent through this client.
144    ///
145    /// Used both by credential providers (token refresh) and by user-initiated
146    /// requests via [`CloudRequestBuilder::send`] and
147    /// [`CloudClient::sign_and_send`]. Override with
148    /// [`CloudClient::with_retry_config`].
149    pub retry_config: RetryConfig,
150    #[cfg(feature = "recording")]
151    recording: Option<RecordingState>,
152}
153
154impl CloudClient {
155    fn new_with_signer(
156        signer: SharedSigner,
157        reqwest_client: Client,
158        service: Arc<dyn HttpService>,
159        retry_config: RetryConfig,
160    ) -> Self {
161        Self {
162            signer,
163            reqwest_client,
164            service,
165            retry_config,
166            #[cfg(feature = "recording")]
167            recording: None,
168        }
169    }
170
171    /// Create a new client with AWS credentials.
172    ///
173    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
174    /// will be spawned on the given runtime handle.
175    pub fn new_aws<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
176    where
177        I: IntoIterator<Item = (K, V)>,
178        K: AsRef<str>,
179        V: Into<String>,
180    {
181        let config = options
182            .into_iter()
183            .fold(
184                aws::AmazonBuilder::new(),
185                |builder, (key, value)| match key.as_ref().parse() {
186                    Ok(k) => builder.with_config(k, value),
187                    Err(_) => builder,
188                },
189            )
190            .build(runtime)?;
191
192        let reqwest_client = config.client_options.client()?;
193        let service = make_service(reqwest_client.clone(), runtime);
194        let retry_config = config.retry_config.clone();
195        Ok(Self::new_with_signer(
196            Arc::new(config),
197            reqwest_client,
198            service,
199            retry_config,
200        ))
201    }
202
203    /// Create a new client with Google Cloud credentials.
204    ///
205    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
206    /// will be spawned on the given runtime handle.
207    pub fn new_google<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
208    where
209        I: IntoIterator<Item = (K, V)>,
210        K: AsRef<str>,
211        V: Into<String>,
212    {
213        let config = options
214            .into_iter()
215            .fold(
216                gcp::GoogleBuilder::new(),
217                |builder, (key, value)| match key.as_ref().parse() {
218                    Ok(k) => builder.with_config(k, value),
219                    Err(_) => builder,
220                },
221            )
222            .build(runtime)?;
223
224        let reqwest_client = config.client_options.client()?;
225        let service = make_service(reqwest_client.clone(), runtime);
226        let retry_config = config.retry_config.clone();
227        Ok(Self::new_with_signer(
228            Arc::new(config),
229            reqwest_client,
230            service,
231            retry_config,
232        ))
233    }
234
235    /// Create a new client with Azure credentials.
236    ///
237    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
238    /// will be spawned on the given runtime handle.
239    pub fn new_azure<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
240    where
241        I: IntoIterator<Item = (K, V)>,
242        K: AsRef<str>,
243        V: Into<String>,
244    {
245        let config = options
246            .into_iter()
247            .fold(
248                azure::AzureBuilder::new(),
249                |builder, (key, value)| match key.as_ref().parse() {
250                    Ok(k) => builder.with_config(k, value),
251                    Err(_) => builder,
252                },
253            )
254            .build(runtime)?;
255
256        let reqwest_client = config.client_options.client()?;
257        let service = make_service(reqwest_client.clone(), runtime);
258        let retry_config = config.retry_config.clone();
259        Ok(Self::new_with_signer(
260            Arc::new(config),
261            reqwest_client,
262            service,
263            retry_config,
264        ))
265    }
266
267    /// Create a new client with Databricks credentials.
268    ///
269    /// Options are applied to a fresh [`databricks::DatabricksBuilder`]; unknown
270    /// keys are ignored. Use [`new_databricks_from_builder`] to start from
271    /// [`databricks::DatabricksBuilder::from_env`] instead, which also loads
272    /// `.databrickscfg` profiles.
273    ///
274    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
275    /// will be spawned on the given runtime handle.
276    ///
277    /// [`new_databricks_from_builder`]: CloudClient::new_databricks_from_builder
278    pub fn new_databricks<I, K, V>(options: I, runtime: Option<&Handle>) -> Result<Self>
279    where
280        I: IntoIterator<Item = (K, V)>,
281        K: AsRef<str>,
282        V: Into<String>,
283    {
284        let builder = options.into_iter().fold(
285            databricks::DatabricksBuilder::new(),
286            |builder, (key, value)| match key.as_ref().parse() {
287                Ok(k) => builder.with_config(k, value),
288                Err(_) => builder,
289            },
290        );
291        Self::new_databricks_from_builder(builder, runtime)
292    }
293
294    /// Create a new client with Databricks credentials from a configured builder.
295    ///
296    /// Unlike [`new_databricks`], this keeps everything the builder already
297    /// resolved — notably the `.databrickscfg` profile values loaded by
298    /// [`databricks::DatabricksBuilder::from_env`].
299    ///
300    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
301    /// will be spawned on the given runtime handle.
302    ///
303    /// [`new_databricks`]: CloudClient::new_databricks
304    pub fn new_databricks_from_builder(
305        builder: databricks::DatabricksBuilder,
306        runtime: Option<&Handle>,
307    ) -> Result<Self> {
308        let config = builder.build(runtime)?;
309
310        let reqwest_client = config.client_options.client()?;
311        let service = make_service(reqwest_client.clone(), runtime);
312        let retry_config = config.retry_config.clone();
313        Ok(Self::new_with_signer(
314            Arc::new(config),
315            reqwest_client,
316            service,
317            retry_config,
318        ))
319    }
320
321    /// Create a new client with a personal access token.
322    pub fn new_with_token(token: impl ToString) -> Self {
323        let reqwest_client = Client::new();
324        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
325        Self::new_with_signer(
326            Arc::new(BearerTokenSigner {
327                token: token.to_string(),
328            }),
329            reqwest_client,
330            service,
331            RetryConfig::default(),
332        )
333    }
334
335    /// Create a new unauthenticated client.
336    pub fn new_unauthenticated() -> Self {
337        let reqwest_client = Client::new();
338        let service: Arc<dyn HttpService> = Arc::new(ReqwestService::new(reqwest_client.clone()));
339        Self::new_with_signer(
340            Arc::new(NoopSigner),
341            reqwest_client,
342            service,
343            RetryConfig::default(),
344        )
345    }
346
347    /// Route all HTTP I/O through the given runtime handle.
348    ///
349    /// This is useful for simple constructors (`new_with_token`, `new_unauthenticated`)
350    /// where no credential providers need to perform HTTP I/O. For cloud provider
351    /// constructors, pass the handle at construction time instead.
352    pub fn with_runtime(mut self, handle: Handle) -> Self {
353        self.service = Arc::new(SpawnService::new(self.service, handle));
354        self
355    }
356
357    /// Replace the [`HttpService`] used for request execution.
358    pub fn with_http_service(mut self, service: Arc<dyn HttpService>) -> Self {
359        self.service = service;
360        self
361    }
362
363    /// Override the [`RetryConfig`] applied to requests sent through this client.
364    ///
365    /// Requests issued via [`CloudRequestBuilder::send`] and
366    /// [`sign_and_send`](Self::sign_and_send) are retried per this config
367    /// (exponential backoff with jitter; safe/idempotent requests are also
368    /// retried on timeout). The [default](RetryConfig::default) allows up to 10
369    /// retries — lower it for latency-sensitive interactive calls.
370    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
371        self.retry_config = retry_config;
372        self
373    }
374
375    /// Sign an already-built [`reqwest::Request`] and dispatch it, with retries.
376    ///
377    /// This is the request-level counterpart to the [`CloudRequestBuilder`]
378    /// flow: it applies the client's [`RequestSigner`] (refreshing credentials
379    /// as needed) and sends the request through the configured [`HttpService`],
380    /// retrying transient failures per the client's [`RetryConfig`]. It is
381    /// useful when the request is produced elsewhere (for example by a protocol
382    /// stack such as ConnectRPC) and only needs authentication and transport.
383    ///
384    /// The request body must be in memory (not a stream): retries clone the
385    /// request, and provider signers such as AWS SigV4 hash the body to sign it.
386    ///
387    /// # Errors
388    ///
389    /// Returns an [`Error`] if signing fails or if the request ultimately fails
390    /// after exhausting retries. A non-success HTTP status code is not itself an
391    /// error; inspect [`reqwest::Response::status`] on the returned response.
392    pub async fn sign_and_send(&self, request: reqwest::Request) -> Result<reqwest::Response> {
393        // The signer operates on a `RequestBuilder`; reconstruct one from the
394        // pre-built request (cloning the in-memory body) so the existing
395        // sign + retry path applies unchanged.
396        let builder = self
397            .reqwest_client
398            .request(request.method().clone(), request.url().clone());
399        let builder = builder.headers(request.headers().clone());
400        let builder = match request.body().and_then(|b| b.as_bytes()) {
401            Some(bytes) => builder.body(bytes.to_vec()),
402            None => builder,
403        };
404        let builder = self.signer.sign(builder).await?;
405        builder
406            .send_retry(&self.retry_config, self.service.clone())
407            .await
408            .map_err(|e| e.error())
409    }
410
411    /// Start building a request for the given HTTP `method` and `url`.
412    ///
413    /// The returned [`CloudRequestBuilder`] borrows this client's signer, so the
414    /// request is signed for the configured provider when
415    /// [`CloudRequestBuilder::send`] is called.
416    pub fn request<U: IntoUrl>(&self, method: Method, url: U) -> CloudRequestBuilder {
417        CloudRequestBuilder {
418            builder: self.reqwest_client.request(method, url),
419            client: self.clone(),
420            #[cfg(feature = "recording")]
421            out_dir: self.recording.as_ref().map(|r| r.out_dir.clone()),
422        }
423    }
424
425    /// Start building a `GET` request for `url`. Shortcut for [`request`](Self::request).
426    pub fn get<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
427        self.request(Method::GET, url)
428    }
429
430    /// Start building a `POST` request for `url`. Shortcut for [`request`](Self::request).
431    pub fn post<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
432        self.request(Method::POST, url)
433    }
434
435    /// Start building a `PUT` request for `url`. Shortcut for [`request`](Self::request).
436    pub fn put<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
437        self.request(Method::PUT, url)
438    }
439
440    /// Start building a `DELETE` request for `url`. Shortcut for [`request`](Self::request).
441    pub fn delete<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
442        self.request(Method::DELETE, url)
443    }
444
445    /// Start building a `HEAD` request for `url`. Shortcut for [`request`](Self::request).
446    pub fn head<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
447        self.request(Method::HEAD, url)
448    }
449
450    /// Start building a `PATCH` request for `url`. Shortcut for [`request`](Self::request).
451    pub fn patch<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
452        self.request(Method::PATCH, url)
453    }
454
455    /// Start building an `OPTIONS` request for `url`. Shortcut for [`request`](Self::request).
456    pub fn options<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
457        self.request(Method::OPTIONS, url)
458    }
459
460    /// Start building a `TRACE` request for `url`. Shortcut for [`request`](Self::request).
461    pub fn trace<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
462        self.request(Method::TRACE, url)
463    }
464
465    /// Start building a `CONNECT` request for `url`. Shortcut for [`request`](Self::request).
466    pub fn connect<U: IntoUrl>(&self, url: U) -> CloudRequestBuilder {
467        self.request(Method::CONNECT, url)
468    }
469
470    /// Enable request/response recording, writing each interaction to `out_dir`.
471    ///
472    /// Once set, every request sent through this client is captured to a
473    /// numbered JSON file (`0000.json`, `0001.json`, …) under `out_dir` for
474    /// later test replay. Sensitive response headers (`authorization`,
475    /// `x-amz-security-token`, `cookie`, and similar) are replaced with
476    /// `"<REDACTED>"` before anything is written to disk, so
477    /// recordings never persist bearer tokens or signing secrets. Request
478    /// headers are not recorded at all.
479    ///
480    /// `out_dir` is canonicalized eagerly, so it must already exist.
481    ///
482    /// Only available when the `recording` feature is enabled.
483    ///
484    /// # Errors
485    ///
486    /// Returns the [`io::Error`](std::io::Error) from canonicalizing `out_dir`,
487    /// for example if the directory does not exist or is not accessible.
488    #[cfg(feature = "recording")]
489    pub fn set_recording_dir(&mut self, out_dir: std::path::PathBuf) -> Result<(), std::io::Error> {
490        let out_dir = std::fs::canonicalize(out_dir)?;
491        self.recording = Some(RecordingState {
492            out_dir,
493            counter: Arc::new(AtomicU64::new(0)),
494        });
495        Ok(())
496    }
497}
498
499/// A builder for a single request issued through a [`CloudClient`].
500///
501/// Created by [`CloudClient::request`] and the per-verb shortcuts such as
502/// [`CloudClient::get`]. Configure the request with the builder methods, then
503/// call [`send`](Self::send) to sign and dispatch it.
504pub struct CloudRequestBuilder {
505    builder: RequestBuilder,
506    client: CloudClient,
507    #[cfg(feature = "recording")]
508    out_dir: Option<PathBuf>,
509}
510
511impl CloudRequestBuilder {
512    /// Add a `Header` to this Request.
513    pub fn header<K, V>(mut self, key: K, value: V) -> CloudRequestBuilder
514    where
515        HeaderName: TryFrom<K>,
516        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
517        HeaderValue: TryFrom<V>,
518        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
519    {
520        self.builder = self.builder.header(key, value);
521        self
522    }
523
524    /// Add a set of Headers to the existing ones on this Request.
525    ///
526    /// The headers will be merged in to any already set.
527    pub fn headers(mut self, headers: HeaderMap) -> CloudRequestBuilder {
528        self.builder = self.builder.headers(headers);
529        self
530    }
531
532    /// Set the request body.
533    pub fn body<T: Into<Body>>(mut self, body: T) -> CloudRequestBuilder {
534        self.builder = self.builder.body(body);
535        self
536    }
537
538    /// Enables a request timeout.
539    ///
540    /// The timeout is applied from when the request starts connecting until the
541    /// response body has finished. It affects only this request and overrides
542    /// the timeout configured using `ClientBuilder::timeout()`.
543    pub fn timeout(mut self, timeout: Duration) -> CloudRequestBuilder {
544        self.builder = self.builder.timeout(timeout);
545        self
546    }
547
548    /// Modify the query string of the URL.
549    ///
550    /// Modifies the URL of this request, adding the parameters provided.
551    /// This method appends and does not overwrite. This means that it can
552    /// be called multiple times and that existing query parameters are not
553    /// overwritten if the same key is used. The key will simply show up
554    /// twice in the query string.
555    /// Calling `.query(&[("foo", "a"), ("foo", "b")])` gives `"foo=a&foo=b"`.
556    ///
557    /// # Note
558    /// This method does not support serializing a single key-value
559    /// pair. Instead of using `.query(("key", "val"))`, use a sequence, such
560    /// as `.query(&[("key", "val")])`. It's also possible to serialize structs
561    /// and maps into a key-value pair.
562    ///
563    /// # Errors
564    /// This method will fail if the object you provide cannot be serialized
565    /// into a query string.
566    pub fn query<T: Serialize + ?Sized>(mut self, query: &T) -> CloudRequestBuilder {
567        self.builder = self.builder.query(query);
568        self
569    }
570
571    /// Send a JSON body.
572    ///
573    /// # Errors
574    ///
575    /// Serialization can fail if `T`'s implementation of `Serialize` decides to
576    /// fail, or if `T` contains a map with non-string keys.
577    pub fn json<T: Serialize + ?Sized>(mut self, json: &T) -> CloudRequestBuilder {
578        self.builder = self.builder.json(json);
579        self
580    }
581
582    /// Sign and send the request, returning the [`reqwest::Response`].
583    ///
584    /// The request is first passed to the client's [`RequestSigner`], which
585    /// attaches the provider-specific authentication (e.g. AWS SigV4 headers or
586    /// an `Authorization: Bearer` header), refreshing any cached credential as
587    /// needed. The signed request is then dispatched through the client's
588    /// [`HttpService`], retrying transient failures per the client's
589    /// [`RetryConfig`] and mapping a non-success status to the matching
590    /// [`Error`].
591    ///
592    /// When the `recording` feature is enabled and a recording directory has
593    /// been configured via `CloudClient::set_recording_dir`, the request and
594    /// response are instead captured to disk (with sensitive headers redacted)
595    /// in a single round-trip — this capture path does **not** retry or map
596    /// status codes to errors, so use the default build for production traffic.
597    ///
598    /// # Errors
599    ///
600    /// Returns an [`Error`] if signing fails, if the request cannot be built, or
601    /// if the underlying HTTP transport returns an error. A non-success HTTP
602    /// status code is not itself an error; inspect [`reqwest::Response::status`]
603    /// on the returned response.
604    pub async fn send(mut self) -> Result<reqwest::Response> {
605        self.builder = self.client.signer.sign(self.builder).await?;
606
607        #[cfg(not(feature = "recording"))]
608        {
609            self.builder
610                .send_retry(&self.client.retry_config, self.client.service.clone())
611                .await
612                .map_err(|e| e.error())
613        }
614        #[cfg(feature = "recording")]
615        {
616            let response = send_record(self).await?;
617            Ok(response)
618        }
619    }
620
621    /// Sign and send the request, preserving the classified non-success result.
622    ///
623    /// This behaves exactly like [`send`](Self::send) — same signing, same
624    /// [`RetryConfig`] semantics (5xx and transport failures retried, 4xx
625    /// returned immediately) — but a non-success outcome is surfaced as the
626    /// [`RetryError`] the retry layer already classified, rather than being
627    /// remapped through [`RetryError::error`] into an opaque [`Error`].
628    ///
629    /// Use this when the caller needs the HTTP **status** and **response body**
630    /// of a failed request — for example to parse a protocol-specific error
631    /// envelope. [`RetryError::status`] and [`RetryError::body`] expose them;
632    /// [`send`](Self::send) discards both by boxing the status into a coarse
633    /// [`Error`] variant (`NotFound`, `AlreadyExists`, …).
634    ///
635    /// When the `recording` feature is enabled the request is captured to disk
636    /// via the same single-shot path as [`send`](Self::send) and the recorded
637    /// [`reqwest::Response`] is returned for any status (no status→error
638    /// mapping), so a non-success status is `Ok(response)` in that build — the
639    /// caller inspects [`reqwest::Response::status`] itself.
640    ///
641    /// # Errors
642    ///
643    /// Returns [`Error`] if signing fails. Otherwise, on a non-success status
644    /// (in the default build) returns [`RetryError`] carrying the status and
645    /// response body; a transport failure surfaces as
646    /// [`RetryError::Reqwest`]/[`RetryError::Transport`].
647    #[cfg(not(feature = "recording"))]
648    pub async fn send_raw(mut self) -> Result<reqwest::Response, SendRawError> {
649        self.builder = self.client.signer.sign(self.builder).await?;
650        self.builder
651            .send_retry(&self.client.retry_config, self.client.service.clone())
652            .await
653            .map_err(SendRawError::Retry)
654    }
655
656    /// See [`send_raw`](Self::send_raw). Under the `recording` feature this maps
657    /// to the single-shot capture path and returns the response for any status.
658    #[cfg(feature = "recording")]
659    pub async fn send_raw(self) -> Result<reqwest::Response, SendRawError> {
660        send_record(self).await
661    }
662}
663
664/// The error returned by [`CloudRequestBuilder::send_raw`].
665///
666/// Distinguishes a failure to *sign* the request (an [`Error`] from the signer,
667/// e.g. a credential-refresh failure) from a failure of the *request itself*
668/// (a [`RetryError`] carrying the HTTP status and response body). This lets a
669/// caller parse a protocol error envelope from [`RetryError::body`] while still
670/// propagating signing failures faithfully.
671#[cfg(not(feature = "recording"))]
672#[derive(Debug, thiserror::Error)]
673pub enum SendRawError {
674    /// The request could not be signed (e.g. credential resolution failed).
675    #[error(transparent)]
676    Sign(#[from] Error),
677    /// The request was sent but failed; carries the classified status + body.
678    #[error(transparent)]
679    Retry(RetryError),
680}
681
682/// Under the `recording` feature `send_raw` returns the recorded response for
683/// any status, so the only error it can produce is a signing/capture [`Error`].
684#[cfg(feature = "recording")]
685pub type SendRawError = Error;
686
687#[cfg(feature = "recording")]
688#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
689pub struct RequestResponseInfo {
690    pub request: RequestInfo,
691    pub response: ResponseInfo,
692}
693
694#[cfg(feature = "recording")]
695#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
696pub struct RequestInfo {
697    pub method: String,
698    pub url_path: String,
699    pub body: Option<String>,
700}
701
702#[cfg(feature = "recording")]
703#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
704pub struct ResponseInfo {
705    pub status: u16,
706    pub headers: HashMap<String, String>,
707    pub body: Option<String>,
708}
709
710/// Header names whose values must never appear in recording files.
711///
712/// These headers carry bearer tokens, signing secrets, or session tokens.
713/// They are replaced with `"<REDACTED>"` before any recording is written to disk.
714#[cfg(feature = "recording")]
715const SENSITIVE_HEADERS: &[&str] = &[
716    "authorization",
717    "x-amz-security-token",
718    "x-amz-content-sha256",
719    "x-databricks-authorization",
720    "x-ms-identity-principal-id",
721    "x-goog-iam-credentials-token",
722    "cookie",
723    "set-cookie",
724];
725
726#[cfg(feature = "recording")]
727fn redact_headers(headers: &HashMap<String, String>) -> HashMap<String, String> {
728    headers
729        .iter()
730        .map(|(k, v)| {
731            let v = if SENSITIVE_HEADERS.contains(&k.to_lowercase().as_str()) {
732                "<REDACTED>".to_string()
733            } else {
734                v.clone()
735            };
736            (k.clone(), v)
737        })
738        .collect()
739}
740
741#[cfg(feature = "recording")]
742async fn send_record(builder: CloudRequestBuilder) -> Result<reqwest::Response> {
743    let Some(out_dir) = builder.out_dir else {
744        let request = builder.builder.build().expect("request to be valid");
745        return builder.client.service.call(request).await;
746    };
747    let (_client, request) = builder.builder.build_split();
748    let request = request.expect("request to be valid");
749
750    let request_info = RequestInfo {
751        method: request.method().as_str().to_string(),
752        url_path: {
753            let url = request.url();
754            match url.query() {
755                Some(query) => format!("{}?{}", url.path(), query),
756                None => url.path().to_string(),
757            }
758        },
759        body: request
760            .body()
761            .and_then(|b| b.as_bytes().map(|b| String::from_utf8_lossy(b).to_string())),
762    };
763
764    let response = builder.client.service.call(request).await?;
765
766    // Record the response
767    let status = response.status().as_u16();
768    let raw_headers: HashMap<String, String> = response
769        .headers()
770        .iter()
771        .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
772        .collect();
773
774    // Get response body while preserving it for the caller
775    let response_bytes = response.bytes().await?;
776    let response_body = if response_bytes.is_empty() {
777        None
778    } else {
779        Some(String::from_utf8_lossy(&response_bytes).to_string())
780    };
781
782    let recording = RequestResponseInfo {
783        request: request_info,
784        response: ResponseInfo {
785            status,
786            // Redact sensitive headers before writing to disk
787            headers: redact_headers(&raw_headers),
788            body: response_body,
789        },
790    };
791
792    let counter = builder
793        .client
794        .recording
795        .as_ref()
796        .map(|r| r.counter.fetch_add(1, Ordering::SeqCst))
797        .unwrap_or(0);
798    let file_path = out_dir.join(format!("{counter:04}.json"));
799    if let Err(e) = std::fs::File::create(&file_path)
800        .and_then(|f| serde_json::to_writer_pretty(f, &recording).map_err(Into::into))
801    {
802        tracing::warn!(
803            "Failed to write recording to {}: {}",
804            file_path.display(),
805            e
806        );
807    }
808
809    // Return a new response built from the recorded data, using the raw (unredacted) headers
810    let mut mock_response = http::Response::builder().status(status);
811    for (k, v) in &raw_headers {
812        mock_response = mock_response.header(k, v);
813    }
814    let mock_response = mock_response
815        .body(response_bytes)
816        .expect("valid status code and headers");
817
818    Ok(reqwest::Response::from(mock_response))
819}
820
821#[cfg(all(test, feature = "recording"))]
822mod tests {
823    use super::*;
824    use std::fs;
825    use tempfile::TempDir;
826
827    #[tokio::test]
828    async fn test_request_response_recording() {
829        // Create a temporary directory for recordings
830        let temp_dir = TempDir::new().unwrap();
831        let temp_path = temp_dir.path().to_path_buf();
832
833        // Set up a mock server
834        let mut server = mockito::Server::new_async().await;
835        let mock = server
836            .mock("GET", "/test")
837            .with_status(200)
838            .with_header("content-type", "application/json")
839            .with_body(r#"{"message": "Hello, World!"}"#)
840            .create_async()
841            .await;
842
843        // Create a cloud client with recording enabled
844        let mut client = CloudClient::new_unauthenticated();
845        client.set_recording_dir(temp_path.clone()).unwrap();
846
847        // Make a request
848        let url = format!("{}/test", server.url());
849        let response = client.get(&url).send().await.unwrap();
850
851        // Verify the response is correct
852        assert_eq!(response.status(), 200);
853        let body = response.text().await.unwrap();
854        assert_eq!(body, r#"{"message": "Hello, World!"}"#);
855
856        // Verify that a recording file was created
857        let recordings: Vec<_> = fs::read_dir(&temp_path)
858            .unwrap()
859            .filter_map(|entry| {
860                let entry = entry.ok()?;
861                let path = entry.path();
862                if path.extension()? == "json" {
863                    Some(path)
864                } else {
865                    None
866                }
867            })
868            .collect();
869
870        assert_eq!(recordings.len(), 1, "Expected exactly one recording file");
871
872        // Read and verify the recording content
873        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
874        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
875
876        // Verify request information
877        assert_eq!(recording.request.method, "GET");
878        assert_eq!(recording.request.url_path, "/test");
879        assert_eq!(recording.request.body, None);
880
881        // Verify response information
882        assert_eq!(recording.response.status, 200);
883        assert_eq!(
884            recording.response.headers.get("content-type").unwrap(),
885            "application/json"
886        );
887        assert_eq!(
888            recording.response.body.as_ref().unwrap(),
889            r#"{"message": "Hello, World!"}"#
890        );
891
892        mock.assert_async().await;
893    }
894
895    #[tokio::test]
896    async fn test_recording_with_request_body() {
897        // Create a temporary directory for recordings
898        let temp_dir = TempDir::new().unwrap();
899        let temp_path = temp_dir.path().to_path_buf();
900
901        // Set up a mock server
902        let mut server = mockito::Server::new_async().await;
903        let mock = server
904            .mock("POST", "/create")
905            .with_status(201)
906            .with_header("location", "/resource/123")
907            .with_body(r#"{"id": 123, "status": "created"}"#)
908            .create_async()
909            .await;
910
911        // Create a cloud client with recording enabled
912        let mut client = CloudClient::new_unauthenticated();
913        client.set_recording_dir(temp_path.clone()).unwrap();
914
915        // Make a POST request with body
916        let url = format!("{}/create", server.url());
917        let response = client
918            .post(&url)
919            .json(&serde_json::json!({"name": "test resource"}))
920            .send()
921            .await
922            .unwrap();
923
924        // Verify the response
925        assert_eq!(response.status(), 201);
926
927        // Verify that a recording file was created
928        let recordings: Vec<_> = fs::read_dir(&temp_path)
929            .unwrap()
930            .filter_map(|entry| {
931                let entry = entry.ok()?;
932                let path = entry.path();
933                if path.extension()? == "json" {
934                    Some(path)
935                } else {
936                    None
937                }
938            })
939            .collect();
940
941        assert_eq!(recordings.len(), 1);
942
943        // Read and verify the recording content
944        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
945        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
946
947        // Verify request information
948        assert_eq!(recording.request.method, "POST");
949        assert_eq!(recording.request.url_path, "/create");
950        assert!(recording.request.body.is_some());
951        assert!(recording.request.body.unwrap().contains("test resource"));
952
953        // Verify response information
954        assert_eq!(recording.response.status, 201);
955        assert_eq!(
956            recording.response.headers.get("location").unwrap(),
957            "/resource/123"
958        );
959        assert!(recording.response.body.unwrap().contains("created"));
960
961        mock.assert_async().await;
962    }
963
964    #[tokio::test]
965    async fn test_counter_based_file_naming() {
966        // Create a temporary directory for recordings
967        let temp_dir = TempDir::new().unwrap();
968        let temp_path = temp_dir.path().to_path_buf();
969
970        // Set up a mock server
971        let mut server = mockito::Server::new_async().await;
972        let mock1 = server
973            .mock("GET", "/first")
974            .with_status(200)
975            .with_body("first response")
976            .create_async()
977            .await;
978        let mock2 = server
979            .mock("GET", "/second")
980            .with_status(200)
981            .with_body("second response")
982            .create_async()
983            .await;
984
985        // Create a cloud client with recording enabled
986        let mut client = CloudClient::new_unauthenticated();
987        client.set_recording_dir(temp_path.clone()).unwrap();
988
989        // Make multiple requests
990        let url1 = format!("{}/first", server.url());
991        let url2 = format!("{}/second", server.url());
992
993        let _response1 = client.get(&url1).send().await.unwrap();
994        let _response2 = client.get(&url2).send().await.unwrap();
995
996        // Verify that files are named with incrementing counter
997        let mut recordings: Vec<_> = fs::read_dir(&temp_path)
998            .unwrap()
999            .filter_map(|entry| {
1000                let entry = entry.ok()?;
1001                let path = entry.path();
1002                if path.extension()? == "json" {
1003                    Some(path)
1004                } else {
1005                    None
1006                }
1007            })
1008            .collect();
1009
1010        recordings.sort();
1011        assert_eq!(recordings.len(), 2);
1012
1013        // Check that files are named 000000.json and 000001.json
1014        assert!(recordings[0].file_name().unwrap().to_str().unwrap() == "0000.json");
1015        assert!(recordings[1].file_name().unwrap().to_str().unwrap() == "0001.json");
1016
1017        // Verify content matches the order of requests
1018        let first_content = fs::read_to_string(&recordings[0]).unwrap();
1019        let first_recording: RequestResponseInfo = serde_json::from_str(&first_content).unwrap();
1020        assert_eq!(first_recording.request.url_path, "/first");
1021        assert_eq!(
1022            first_recording.response.body.as_ref().unwrap(),
1023            "first response"
1024        );
1025
1026        let second_content = fs::read_to_string(&recordings[1]).unwrap();
1027        let second_recording: RequestResponseInfo = serde_json::from_str(&second_content).unwrap();
1028        assert_eq!(second_recording.request.url_path, "/second");
1029        assert_eq!(
1030            second_recording.response.body.as_ref().unwrap(),
1031            "second response"
1032        );
1033
1034        mock1.assert_async().await;
1035        mock2.assert_async().await;
1036    }
1037
1038    #[tokio::test]
1039    async fn test_query_parameter_recording() {
1040        // Create a temporary directory for recordings
1041        let temp_dir = TempDir::new().unwrap();
1042        let temp_path = temp_dir.path().to_path_buf();
1043
1044        // Start a mock server
1045        let mut server = mockito::Server::new_async().await;
1046
1047        // Create a mock that expects query parameters
1048        let mock = server
1049            .mock("GET", "/catalogs?max_results=10&page_token=abc123")
1050            .with_status(200)
1051            .with_header("content-type", "application/json")
1052            .with_body(r#"{"catalogs": []}"#)
1053            .create_async()
1054            .await;
1055
1056        // Create a client with recording enabled
1057        let mut client = CloudClient::new_unauthenticated();
1058        client.set_recording_dir(temp_path.clone()).unwrap();
1059
1060        // Make a request with query parameters
1061        let url = format!("{}/catalogs?max_results=10&page_token=abc123", server.url());
1062        let response = client.get(&url).send().await.unwrap();
1063
1064        assert!(response.status().is_success());
1065
1066        // Verify that the recording file was created
1067        let recordings: Vec<_> = fs::read_dir(&temp_path)
1068            .unwrap()
1069            .filter_map(|entry| {
1070                let entry = entry.ok()?;
1071                let path = entry.path();
1072                if path.extension()? == "json" {
1073                    Some(path)
1074                } else {
1075                    None
1076                }
1077            })
1078            .collect();
1079
1080        assert_eq!(recordings.len(), 1);
1081
1082        // Read and verify the recording content includes query parameters
1083        let recording_content = fs::read_to_string(&recordings[0]).unwrap();
1084        let recording: RequestResponseInfo = serde_json::from_str(&recording_content).unwrap();
1085
1086        // Verify request information includes query parameters
1087        assert_eq!(recording.request.method, "GET");
1088        assert_eq!(
1089            recording.request.url_path,
1090            "/catalogs?max_results=10&page_token=abc123"
1091        );
1092        assert_eq!(recording.request.body, None);
1093
1094        // Verify response information
1095        assert_eq!(recording.response.status, 200);
1096        assert_eq!(
1097            recording.response.body.as_ref().unwrap(),
1098            r#"{"catalogs": []}"#
1099        );
1100
1101        mock.assert_async().await;
1102    }
1103
1104    /// Verify that the bearer token injected by a signed request does not appear
1105    /// in the recording file. Request headers are currently not recorded, so this
1106    /// test confirms the token does not leak via any other path (e.g. echoed back
1107    /// in a response header or response body).
1108    #[tokio::test]
1109    async fn test_recording_does_not_contain_bearer_token_value() {
1110        let temp_dir = TempDir::new().unwrap();
1111        let temp_path = temp_dir.path().to_path_buf();
1112
1113        let mut server = mockito::Server::new_async().await;
1114        let mock = server
1115            .mock("GET", "/secret")
1116            .match_header("authorization", mockito::Matcher::Any)
1117            .with_status(200)
1118            // Server does NOT echo the token back — simulates a well-behaved API
1119            .with_body("ok")
1120            .create_async()
1121            .await;
1122
1123        let mut client = CloudClient::new_with_token("super-secret-token-12345");
1124        client.set_recording_dir(temp_path.clone()).unwrap();
1125
1126        let url = format!("{}/secret", server.url());
1127        client.get(&url).send().await.unwrap();
1128
1129        let recording_path = temp_path.join("0000.json");
1130        let content = fs::read_to_string(&recording_path).unwrap();
1131
1132        // The raw token must not appear in the file at all
1133        assert!(
1134            !content.contains("super-secret-token-12345"),
1135            "raw bearer token leaked into recording: {content}"
1136        );
1137
1138        mock.assert_async().await;
1139    }
1140
1141    /// Verify that sensitive headers returned by the server (e.g. a reflected
1142    /// Authorization or AWS security token) are redacted before being written
1143    /// to the recording file.
1144    #[tokio::test]
1145    async fn test_recording_redacts_sensitive_response_headers() {
1146        let temp_dir = TempDir::new().unwrap();
1147        let temp_path = temp_dir.path().to_path_buf();
1148
1149        let mut server = mockito::Server::new_async().await;
1150        // Simulate a server that echoes a sensitive header back in its response
1151        let mock = server
1152            .mock("GET", "/s3")
1153            .with_status(200)
1154            .with_header("x-amz-security-token", "AQoXnyc4LLI2AJvUAMOGAR8a1234567890")
1155            .with_header("authorization", "Bearer should-be-redacted")
1156            .with_body("{}")
1157            .create_async()
1158            .await;
1159
1160        let mut client = CloudClient::new_unauthenticated();
1161        client.set_recording_dir(temp_path.clone()).unwrap();
1162
1163        let url = format!("{}/s3", server.url());
1164        client.get(&url).send().await.unwrap();
1165
1166        let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1167        assert!(
1168            !content.contains("AQoXnyc4LLI2AJvUAMOGAR8a1234567890"),
1169            "x-amz-security-token leaked into recording: {content}"
1170        );
1171        assert!(
1172            !content.contains("should-be-redacted"),
1173            "Authorization value leaked into recording: {content}"
1174        );
1175        // Both headers should be replaced with the redaction sentinel
1176        assert_eq!(
1177            content.matches("<REDACTED>").count(),
1178            2,
1179            "expected 2 <REDACTED> entries in recording: {content}"
1180        );
1181
1182        mock.assert_async().await;
1183    }
1184
1185    /// Verify that a recording produced with redacted headers can still be parsed.
1186    #[tokio::test]
1187    async fn test_recording_remains_valid_json_after_redaction() {
1188        let temp_dir = TempDir::new().unwrap();
1189        let temp_path = temp_dir.path().to_path_buf();
1190
1191        let mut server = mockito::Server::new_async().await;
1192        let _mock = server
1193            .mock("GET", "/check")
1194            .with_status(200)
1195            .with_header("authorization", "Bearer should-be-redacted")
1196            .with_body(r#"{"ok":true}"#)
1197            .create_async()
1198            .await;
1199
1200        let mut client = CloudClient::new_unauthenticated();
1201        client.set_recording_dir(temp_path.clone()).unwrap();
1202
1203        let url = format!("{}/check", server.url());
1204        client.get(&url).send().await.unwrap();
1205
1206        let content = fs::read_to_string(temp_path.join("0000.json")).unwrap();
1207        // Must parse cleanly as a RequestResponseInfo
1208        let parsed: RequestResponseInfo = serde_json::from_str(&content)
1209            .expect("recording file must be valid JSON even after redaction");
1210        assert_eq!(parsed.response.status, 200);
1211    }
1212}
1213
1214// These assert the production send path's retry + error-mapping contract.
1215// The `recording` feature replaces that path with a single-shot capture (no
1216// retry, no status->error mapping — see `send_record`), so they are scoped to
1217// the non-recording build where the behavior under test actually applies.
1218#[cfg(all(test, not(feature = "recording")))]
1219mod retry_integration_tests {
1220    use super::*;
1221    use std::time::Duration;
1222
1223    fn fast_retry(max_retries: usize) -> RetryConfig {
1224        RetryConfig {
1225            backoff: crate::backoff::BackoffConfig {
1226                init_backoff: Duration::from_millis(1),
1227                max_backoff: Duration::from_millis(5),
1228                base: 2.,
1229            },
1230            max_retries,
1231            retry_timeout: Duration::from_secs(30),
1232        }
1233    }
1234
1235    // A 5xx is retried by CloudClient::send: a single 503 followed by a 200
1236    // should surface the 200, proving the user-request path now honors
1237    // retry_config (previously it did a bare service.call with no retry).
1238    #[tokio::test]
1239    async fn send_retries_server_error_then_succeeds() {
1240        let mut server = mockito::Server::new_async().await;
1241        let fail = server
1242            .mock("GET", "/r")
1243            .with_status(503)
1244            .expect(1)
1245            .create_async()
1246            .await;
1247        let ok = server
1248            .mock("GET", "/r")
1249            .with_status(200)
1250            .with_body("ok")
1251            .expect(1)
1252            .create_async()
1253            .await;
1254
1255        let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1256        let resp = client
1257            .get(format!("{}/r", server.url()))
1258            .send()
1259            .await
1260            .unwrap();
1261
1262        assert_eq!(resp.status(), 200);
1263        assert_eq!(resp.text().await.unwrap(), "ok");
1264        fail.assert_async().await;
1265        ok.assert_async().await;
1266    }
1267
1268    // A 4xx is not retryable: it must surface immediately as an error without
1269    // consuming retries (a second mock would go unmatched).
1270    #[tokio::test]
1271    async fn send_does_not_retry_client_error() {
1272        let mut server = mockito::Server::new_async().await;
1273        let mock = server
1274            .mock("GET", "/r")
1275            .with_status(404)
1276            .expect(1)
1277            .create_async()
1278            .await;
1279
1280        let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1281        let err = client
1282            .get(format!("{}/r", server.url()))
1283            .send()
1284            .await
1285            .unwrap_err();
1286
1287        assert!(matches!(err, Error::NotFound { .. }), "got {err:?}");
1288        mock.assert_async().await;
1289    }
1290
1291    // send_raw preserves the classified status + body of a non-success
1292    // response instead of remapping it into an opaque Error. A 404 and a 409
1293    // must both surface as Err(RetryError) whose status()/body() are intact,
1294    // and a 5xx must still be retried (503 then 200 -> Ok(200)).
1295    #[tokio::test]
1296    async fn send_raw_preserves_status_and_body() {
1297        let mut server = mockito::Server::new_async().await;
1298
1299        // 404 with a body envelope: status + body must survive.
1300        let not_found = server
1301            .mock("GET", "/nf")
1302            .with_status(404)
1303            .with_body(r#"{"error":{"type":"NoSuchTableException"}}"#)
1304            .expect(1)
1305            .create_async()
1306            .await;
1307        // 409 conflict: same — non-retryable, body preserved.
1308        let conflict = server
1309            .mock("POST", "/cf")
1310            .with_status(409)
1311            .with_body(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1312            .expect(1)
1313            .create_async()
1314            .await;
1315        // 5xx is still retried under send_raw: one 503 then a 200.
1316        let fail = server
1317            .mock("GET", "/rt")
1318            .with_status(503)
1319            .expect(1)
1320            .create_async()
1321            .await;
1322        let ok = server
1323            .mock("GET", "/rt")
1324            .with_status(200)
1325            .with_body("ok")
1326            .expect(1)
1327            .create_async()
1328            .await;
1329
1330        let client = CloudClient::new_unauthenticated().with_retry_config(fast_retry(3));
1331
1332        let err = client
1333            .get(format!("{}/nf", server.url()))
1334            .send_raw()
1335            .await
1336            .unwrap_err();
1337        match err {
1338            SendRawError::Retry(e) => {
1339                assert_eq!(e.status(), Some(reqwest::StatusCode::NOT_FOUND));
1340                assert_eq!(
1341                    e.body(),
1342                    Some(r#"{"error":{"type":"NoSuchTableException"}}"#)
1343                );
1344            }
1345            other => panic!("expected Retry error, got {other:?}"),
1346        }
1347
1348        let err = client
1349            .post(format!("{}/cf", server.url()))
1350            .send_raw()
1351            .await
1352            .unwrap_err();
1353        match err {
1354            SendRawError::Retry(e) => {
1355                assert_eq!(e.status(), Some(reqwest::StatusCode::CONFLICT));
1356                assert_eq!(
1357                    e.body(),
1358                    Some(r#"{"error":{"type":"CommitVersionConflictException"}}"#)
1359                );
1360            }
1361            other => panic!("expected Retry error, got {other:?}"),
1362        }
1363
1364        let resp = client
1365            .get(format!("{}/rt", server.url()))
1366            .send_raw()
1367            .await
1368            .unwrap();
1369        assert_eq!(resp.status(), 200);
1370        assert_eq!(resp.text().await.unwrap(), "ok");
1371
1372        not_found.assert_async().await;
1373        conflict.assert_async().await;
1374        fail.assert_async().await;
1375        ok.assert_async().await;
1376    }
1377
1378    // sign_and_send takes a pre-built reqwest::Request, applies the signer
1379    // (here a bearer token), and retries transient failures just like send.
1380    #[tokio::test]
1381    async fn sign_and_send_signs_and_retries() {
1382        let mut server = mockito::Server::new_async().await;
1383        let fail = server
1384            .mock("POST", "/rpc")
1385            .match_header("authorization", "Bearer tok")
1386            .with_status(503)
1387            .expect(1)
1388            .create_async()
1389            .await;
1390        let ok = server
1391            .mock("POST", "/rpc")
1392            .match_header("authorization", "Bearer tok")
1393            .with_status(200)
1394            .with_body("pong")
1395            .expect(1)
1396            .create_async()
1397            .await;
1398
1399        let client = CloudClient::new_with_token("tok").with_retry_config(fast_retry(3));
1400        let request = reqwest::Client::new()
1401            .post(format!("{}/rpc", server.url()))
1402            .body("ping")
1403            .build()
1404            .unwrap();
1405        let resp = client.sign_and_send(request).await.unwrap();
1406
1407        assert_eq!(resp.status(), 200);
1408        assert_eq!(resp.text().await.unwrap(), "pong");
1409        fail.assert_async().await;
1410        ok.assert_async().await;
1411    }
1412}
1413
1414#[cfg(test)]
1415mod databricks_client_tests {
1416    use super::*;
1417    use databricks::{DatabricksBuilder, DatabricksConfigKey};
1418
1419    /// A pre-configured builder is used as-is: its credential signs requests.
1420    #[tokio::test]
1421    async fn new_databricks_from_builder_signs_with_builder_credential() {
1422        let mut server = mockito::Server::new_async().await;
1423        let mock = server
1424            .mock("GET", "/api/2.1/unity-catalog/catalogs")
1425            .match_header("authorization", "Bearer builder-pat")
1426            .with_status(200)
1427            .create_async()
1428            .await;
1429
1430        let builder = DatabricksBuilder::new()
1431            .with_config(DatabricksConfigKey::Host, server.url())
1432            .with_config(DatabricksConfigKey::Token, "builder-pat")
1433            .with_config(
1434                DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1435                "true",
1436            );
1437        let client = CloudClient::new_databricks_from_builder(builder, None).unwrap();
1438
1439        let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1440        let response = client.get(&url).send().await.unwrap();
1441        assert_eq!(response.status(), 200);
1442        mock.assert_async().await;
1443    }
1444
1445    /// The option-list constructor still parses keys (and ignores unknown ones).
1446    #[tokio::test]
1447    async fn new_databricks_applies_string_options() {
1448        let mut server = mockito::Server::new_async().await;
1449        let mock = server
1450            .mock("GET", "/ping")
1451            .match_header("authorization", "Bearer option-pat")
1452            .with_status(200)
1453            .create_async()
1454            .await;
1455
1456        let client = CloudClient::new_databricks(
1457            [
1458                ("databricks_host", server.url()),
1459                ("databricks_token", "option-pat".to_string()),
1460                ("allow_http", "true".to_string()),
1461                ("not_a_real_key", "ignored".to_string()),
1462            ],
1463            None,
1464        )
1465        .unwrap();
1466
1467        let response = client
1468            .get(format!("{}/ping", server.url()))
1469            .send()
1470            .await
1471            .unwrap();
1472        assert_eq!(response.status(), 200);
1473        mock.assert_async().await;
1474    }
1475}
1476
1477#[cfg(test)]
1478mod databricks_azure_auth_tests {
1479    use super::*;
1480    use databricks::{DatabricksBuilder, DatabricksConfigKey};
1481
1482    fn azure_sp_builder(server: &mockito::Server) -> DatabricksBuilder {
1483        DatabricksBuilder::new()
1484            .with_config(DatabricksConfigKey::Host, server.url())
1485            .with_config(DatabricksConfigKey::ClientId, "sp-id")
1486            .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1487            .with_config(DatabricksConfigKey::AzureTenantId, "tenant-1")
1488            .with_config(DatabricksConfigKey::AzureAuthorityHost, server.url())
1489            .with_config(
1490                DatabricksConfigKey::Client(ClientConfigKey::AllowHttp),
1491                "true",
1492            )
1493    }
1494
1495    /// A tenant id selects Azure AD client credentials for the Azure
1496    /// Databricks application instead of Databricks OAuth M2M.
1497    #[tokio::test]
1498    async fn tenant_selects_azure_client_secret_auth() {
1499        let mut server = mockito::Server::new_async().await;
1500        let aad = server
1501            .mock("POST", "/tenant-1/oauth2/v2.0/token")
1502            .match_body(mockito::Matcher::AllOf(vec![
1503                mockito::Matcher::UrlEncoded("client_id".into(), "sp-id".into()),
1504                mockito::Matcher::UrlEncoded(
1505                    "scope".into(),
1506                    "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default".into(),
1507                ),
1508            ]))
1509            .with_status(200)
1510            .with_header("content-type", "application/json")
1511            .with_body(r#"{"access_token":"aad-token","expires_in":3600,"token_type":"Bearer"}"#)
1512            .expect(1)
1513            .create_async()
1514            .await;
1515        let m2m = server
1516            .mock("POST", "/oidc/v1/token")
1517            .expect(0)
1518            .create_async()
1519            .await;
1520        let api = server
1521            .mock("GET", "/api/2.1/unity-catalog/catalogs")
1522            .match_header("authorization", "Bearer aad-token")
1523            .with_status(200)
1524            .create_async()
1525            .await;
1526
1527        let client =
1528            CloudClient::new_databricks_from_builder(azure_sp_builder(&server), None).unwrap();
1529        let url = format!("{}/api/2.1/unity-catalog/catalogs", server.url());
1530        assert_eq!(client.get(&url).send().await.unwrap().status(), 200);
1531
1532        aad.assert_async().await;
1533        m2m.assert_async().await;
1534        api.assert_async().await;
1535    }
1536
1537    #[test]
1538    fn forced_azure_client_secret_requires_a_tenant() {
1539        let result = DatabricksBuilder::new()
1540            .with_config(DatabricksConfigKey::ClientId, "sp-id")
1541            .with_config(DatabricksConfigKey::ClientSecret, "sp-secret")
1542            .with_auth_type("azure-client-secret")
1543            .build(None);
1544        let Err(err) = result else {
1545            panic!("expected an error without a tenant id")
1546        };
1547        assert!(err.to_string().contains("AZURE_TENANT_ID"), "{err}");
1548    }
1549
1550    #[test]
1551    fn azure_config_keys_roundtrip() {
1552        for key in [
1553            DatabricksConfigKey::AzureTenantId,
1554            DatabricksConfigKey::AzureAuthorityHost,
1555        ] {
1556            assert_eq!(key.as_ref().parse::<DatabricksConfigKey>().unwrap(), key);
1557        }
1558        assert_eq!(
1559            "azure_tenant_id".parse::<DatabricksConfigKey>().unwrap(),
1560            DatabricksConfigKey::AzureTenantId
1561        );
1562    }
1563}