tracedb_sdk/core/
http_client.rs

1use crate::{join_url, ApiError, ClientConfig, OAuthTokenProvider, RequestOptions};
2use futures::{future::BoxFuture, Stream, StreamExt};
3use reqwest::{
4    header::{HeaderMap, HeaderName, HeaderValue},
5    Client, Method, Request, Response,
6};
7use serde::de::DeserializeOwned;
8
9use std::{
10    pin::Pin,
11    str::FromStr,
12    sync::Arc,
13    task::{Context, Poll},
14};
15
16/// A parsed HTTP response that includes the deserialized body along with
17/// the HTTP status code and response headers.
18#[derive(Debug)]
19pub struct RawResponse<T> {
20    /// The deserialized response body.
21    pub body: T,
22    /// The HTTP status code of the response.
23    pub status_code: u16,
24    /// The HTTP response headers.
25    pub headers: HeaderMap,
26}
27
28/// A streaming byte stream for downloading files efficiently
29pub struct ByteStream {
30    content_length: Option<u64>,
31    inner: Pin<Box<dyn Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send>>,
32}
33
34impl ByteStream {
35    /// Create a new ByteStream from a Response
36    pub(crate) fn new(response: Response) -> Self {
37        let content_length = response.content_length();
38        let stream = response.bytes_stream();
39
40        Self {
41            content_length,
42            inner: Box::pin(stream),
43        }
44    }
45
46    /// Collect the entire stream into a `Vec<u8>`
47    ///
48    /// This consumes the stream and buffers all data into memory.
49    /// For large files, prefer using `try_next()` to process chunks incrementally.
50    ///
51    /// # Example
52    /// ```no_run
53    /// let stream = client.download_file().await?;
54    /// let bytes = stream.collect().await?;
55    /// ```
56    pub async fn collect(mut self) -> Result<Vec<u8>, ApiError> {
57        let mut result = Vec::new();
58        while let Some(chunk) = self.inner.next().await {
59            result.extend_from_slice(&chunk.map_err(ApiError::Network)?);
60        }
61        Ok(result)
62    }
63
64    /// Get the next chunk from the stream
65    ///
66    /// Returns `Ok(Some(bytes))` if a chunk is available,
67    /// `Ok(None)` if the stream is finished, or an error.
68    ///
69    /// # Example
70    /// ```no_run
71    /// let mut stream = client.download_file().await?;
72    /// while let Some(chunk) = stream.try_next().await? {
73    ///     process_chunk(&chunk);
74    /// }
75    /// ```
76    pub async fn try_next(&mut self) -> Result<Option<bytes::Bytes>, ApiError> {
77        match self.inner.next().await {
78            Some(Ok(bytes)) => Ok(Some(bytes)),
79            Some(Err(e)) => Err(ApiError::Network(e)),
80            None => Ok(None),
81        }
82    }
83
84    /// Get the content length from response headers if available
85    pub fn content_length(&self) -> Option<u64> {
86        self.content_length
87    }
88}
89
90impl Stream for ByteStream {
91    type Item = Result<bytes::Bytes, ApiError>;
92
93    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
94        match self.inner.as_mut().poll_next(cx) {
95            Poll::Ready(Some(Ok(bytes))) => Poll::Ready(Some(Ok(bytes))),
96            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(ApiError::Network(e)))),
97            Poll::Ready(None) => Poll::Ready(None),
98            Poll::Pending => Poll::Pending,
99        }
100    }
101}
102
103/// Trait for executing HTTP requests, enabling injection of custom
104/// transport implementations (e.g., for CLI execution-sharing).
105///
106/// When an external executor is provided, the SDK delegates raw HTTP
107/// execution to it, allowing the caller's transport stack to handle
108/// auth, retries, and TLS configuration.
109#[doc(hidden)]
110pub trait RequestExecutor: Send + Sync {
111    fn execute(
112        &self,
113        request: Request,
114    ) -> BoxFuture<'_, Result<Response, Box<dyn std::error::Error + Send + Sync>>>;
115}
116
117/// Default executor that delegates to a `reqwest::Client`.
118struct ReqwestExecutor {
119    client: Client,
120}
121
122impl RequestExecutor for ReqwestExecutor {
123    fn execute(
124        &self,
125        request: Request,
126    ) -> BoxFuture<'_, Result<Response, Box<dyn std::error::Error + Send + Sync>>> {
127        Box::pin(async move {
128            self.client
129                .execute(request)
130                .await
131                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
132        })
133    }
134}
135
136/// Configuration for OAuth token fetching.
137///
138/// This struct contains all the information needed to automatically fetch
139/// and refresh OAuth tokens.
140#[derive(Clone)]
141pub struct OAuthConfig {
142    /// The OAuth token provider that manages token caching and refresh
143    pub token_provider: Arc<OAuthTokenProvider>,
144    /// The token endpoint path (e.g., "/token")
145    pub token_endpoint: String,
146}
147
148/// Response from an OAuth token endpoint.
149#[derive(Debug, Clone, serde::Deserialize)]
150struct OAuthTokenResponse {
151    access_token: String,
152    #[serde(default)]
153    expires_in: Option<i64>,
154}
155
156/// Internal HTTP client that handles requests with authentication and retries
157#[derive(Clone)]
158pub struct HttpClient {
159    client: Client,
160    executor: Option<Arc<dyn RequestExecutor>>,
161    config: ClientConfig,
162    /// Optional OAuth configuration for automatic token management
163    oauth_config: Option<OAuthConfig>,
164}
165
166impl HttpClient {
167    /// Creates a new HttpClient without OAuth support.
168    pub fn new(config: ClientConfig) -> Result<Self, ApiError> {
169        Self::new_with_oauth(config, None)
170    }
171
172    /// Creates a new HttpClient with optional OAuth support.
173    ///
174    /// When `oauth_config` is provided, the client will automatically fetch and refresh
175    /// OAuth tokens before making requests.
176    pub fn new_with_oauth(
177        config: ClientConfig,
178        oauth_config: Option<OAuthConfig>,
179    ) -> Result<Self, ApiError> {
180        let client = Client::builder()
181            .timeout(config.timeout)
182            .user_agent(&config.user_agent)
183            .build()
184            .map_err(ApiError::Network)?;
185
186        Ok(Self {
187            client,
188            executor: None,
189            config,
190            oauth_config,
191        })
192    }
193
194    /// Creates an HttpClient with an injected request executor.
195    ///
196    /// When using an injected executor, the client delegates HTTP execution
197    /// entirely to the executor. Auth headers, custom headers, and retry
198    /// logic are NOT applied by this client — the executor's transport
199    /// stack is expected to handle them. This prevents double-retry and
200    /// double-auth when the SDK is embedded inside a CLI.
201    #[doc(hidden)]
202    pub fn with_executor(executor: Arc<dyn RequestExecutor>, config: ClientConfig) -> Self {
203        let client = Client::new();
204        Self {
205            client,
206            executor: Some(executor),
207            config,
208            oauth_config: None,
209        }
210    }
211
212    /// Returns the configured base URL.
213    pub fn base_url(&self) -> &str {
214        &self.config.base_url
215    }
216
217    /// Returns a reference to the client configuration.
218    pub fn config(&self) -> &ClientConfig {
219        &self.config
220    }
221
222    /// Execute a request and return the parsed body along with HTTP status code and headers.
223    ///
224    /// Unlike `execute_request`, this method preserves the HTTP metadata from the response,
225    /// which is useful for paginated endpoints where callers need access to status codes
226    /// and headers alongside the deserialized body.
227    pub async fn execute_request_raw<T>(
228        &self,
229        method: Method,
230        path: &str,
231        body: Option<serde_json::Value>,
232        query_params: Option<Vec<(String, String)>>,
233        options: Option<RequestOptions>,
234    ) -> Result<RawResponse<T>, ApiError>
235    where
236        T: DeserializeOwned,
237    {
238        let url = join_url(&self.config.base_url, path);
239        let mut request = self.client.request(method, &url);
240
241        if let Some(params) = query_params {
242            request = request.query(&params);
243        }
244
245        if let Some(opts) = &options {
246            if !opts.additional_query_params.is_empty() {
247                request = request.query(&opts.additional_query_params);
248            }
249        }
250
251        if let Some(body) = body {
252            request = request.json(&body);
253        }
254
255        let req = request.build().map_err(|e| ApiError::Network(e))?;
256
257        let response = self.send_request(req, &options).await?;
258        self.parse_response_raw(response).await
259    }
260
261    /// Execute a request with the given method, path, and options
262    pub async fn execute_request<T>(
263        &self,
264        method: Method,
265        path: &str,
266        body: Option<serde_json::Value>,
267        query_params: Option<Vec<(String, String)>>,
268        options: Option<RequestOptions>,
269    ) -> Result<T, ApiError>
270    where
271        T: DeserializeOwned,
272    {
273        let url = join_url(&self.config.base_url, path);
274        let mut request = self.client.request(method, &url);
275
276        if let Some(params) = query_params {
277            request = request.query(&params);
278        }
279
280        if let Some(opts) = &options {
281            if !opts.additional_query_params.is_empty() {
282                request = request.query(&opts.additional_query_params);
283            }
284        }
285
286        if let Some(body) = body {
287            request = request.json(&body);
288        }
289
290        let req = request.build().map_err(|e| ApiError::Network(e))?;
291
292        let response = self.send_request(req, &options).await?;
293        self.parse_response(response).await
294    }
295
296    /// Execute a request with an explicit base URL override.
297    ///
298    /// Used for multi-URL environments where different endpoints
299    /// resolve to different base URLs.
300    pub async fn execute_request_with_base_url<T>(
301        &self,
302        base_url: &str,
303        method: Method,
304        path: &str,
305        body: Option<serde_json::Value>,
306        query_params: Option<Vec<(String, String)>>,
307        options: Option<RequestOptions>,
308    ) -> Result<T, ApiError>
309    where
310        T: DeserializeOwned,
311    {
312        let url = join_url(base_url, path);
313        let mut request = self.client.request(method, &url);
314
315        if let Some(params) = query_params {
316            request = request.query(&params);
317        }
318
319        if let Some(opts) = &options {
320            if !opts.additional_query_params.is_empty() {
321                request = request.query(&opts.additional_query_params);
322            }
323        }
324
325        if let Some(body) = body {
326            request = request.json(&body);
327        }
328
329        let req = request.build().map_err(|e| ApiError::Network(e))?;
330
331        let response = self.send_request(req, &options).await?;
332        self.parse_response(response).await
333    }
334
335    /// Applies auth/headers and executes the request, choosing between
336    /// the injected executor path (no SDK-level auth/headers/retries)
337    /// and the default path (full SDK behavior).
338    async fn send_request(
339        &self,
340        req: Request,
341        options: &Option<RequestOptions>,
342    ) -> Result<Response, ApiError> {
343        if let Some(executor) = &self.executor {
344            executor.execute(req).await.map_err(ApiError::Executor)
345        } else {
346            let mut req = req;
347            self.apply_auth_headers(&mut req, options).await?;
348            self.apply_custom_headers(&mut req, options)?;
349            self.execute_with_retries(req, options).await
350        }
351    }
352
353    async fn apply_auth_headers(
354        &self,
355        request: &mut Request,
356        options: &Option<RequestOptions>,
357    ) -> Result<(), ApiError> {
358        let headers = request.headers_mut();
359
360        // Apply API key (request options override config)
361        let api_key = options
362            .as_ref()
363            .and_then(|opts| opts.api_key.as_ref())
364            .or(self.config.api_key.as_ref());
365
366        if let Some(key) = api_key {
367            let header_value = key.to_string();
368            headers.insert(
369                "api_key",
370                header_value.parse().map_err(|_| ApiError::InvalidHeader)?,
371            );
372        }
373
374        // Apply bearer token - priority: request options > OAuth > config
375        let token = if let Some(opts) = options.as_ref() {
376            if opts.token.is_some() {
377                opts.token.clone()
378            } else {
379                None
380            }
381        } else {
382            None
383        };
384
385        let token = match token {
386            Some(t) => Some(t),
387            None => {
388                // Try OAuth token provider if configured
389                if let Some(oauth_config) = &self.oauth_config {
390                    Some(self.get_oauth_token(oauth_config).await?)
391                } else {
392                    // Fall back to static token from config
393                    self.config.token.clone()
394                }
395            }
396        };
397
398        if let Some(token) = token {
399            let auth_value = format!("Bearer {}", token);
400            headers.insert(
401                "Authorization",
402                auth_value.parse().map_err(|_| ApiError::InvalidHeader)?,
403            );
404        }
405
406        Ok(())
407    }
408
409    /// Fetches an OAuth token, using the cached token if valid or fetching a new one.
410    async fn get_oauth_token(&self, oauth_config: &OAuthConfig) -> Result<String, ApiError> {
411        let token_provider = &oauth_config.token_provider;
412        let token_endpoint = &oauth_config.token_endpoint;
413        let client_id = token_provider.client_id().to_string();
414        let client_secret = token_provider.client_secret().to_string();
415        let base_url = self.config.base_url.clone();
416
417        // Use the async get_or_fetch method with a closure that fetches the token
418        token_provider
419            .get_or_fetch_async(|| async {
420                self.fetch_oauth_token(&base_url, token_endpoint, &client_id, &client_secret)
421                    .await
422            })
423            .await
424    }
425
426    /// Makes an HTTP request to the OAuth token endpoint to fetch a new token.
427    async fn fetch_oauth_token(
428        &self,
429        base_url: &str,
430        token_endpoint: &str,
431        client_id: &str,
432        client_secret: &str,
433    ) -> Result<(String, u64), ApiError> {
434        let url = join_url(base_url, token_endpoint);
435
436        // Build the token request body
437        let body = serde_json::json!({
438            "client_id": client_id,
439            "client_secret": client_secret,
440            "grant_type": "client_credentials"
441        });
442
443        let response = self
444            .client
445            .request(Method::POST, &url)
446            .json(&body)
447            .send()
448            .await
449            .map_err(ApiError::Network)?;
450
451        if !response.status().is_success() {
452            let status_code = response.status().as_u16();
453            let body = response.text().await.ok();
454            return Err(ApiError::from_response(status_code, body.as_deref()));
455        }
456
457        // Parse the token response
458        let token_response: OAuthTokenResponse =
459            response.json().await.map_err(ApiError::Network)?;
460
461        let expires_in = token_response.expires_in.unwrap_or(3600) as u64;
462        Ok((token_response.access_token, expires_in))
463    }
464
465    fn apply_custom_headers(
466        &self,
467        request: &mut Request,
468        options: &Option<RequestOptions>,
469    ) -> Result<(), ApiError> {
470        let headers = request.headers_mut();
471
472        // Apply config-level custom headers
473        for (key, value) in &self.config.custom_headers {
474            headers.insert(
475                HeaderName::from_str(key).map_err(|_| ApiError::InvalidHeader)?,
476                HeaderValue::from_str(value).map_err(|_| ApiError::InvalidHeader)?,
477            );
478        }
479
480        // Apply request-level custom headers (override config)
481        if let Some(options) = options {
482            for (key, value) in &options.additional_headers {
483                headers.insert(
484                    HeaderName::from_str(key).map_err(|_| ApiError::InvalidHeader)?,
485                    HeaderValue::from_str(value).map_err(|_| ApiError::InvalidHeader)?,
486                );
487            }
488        }
489
490        Ok(())
491    }
492
493    async fn execute_with_retries(
494        &self,
495        request: Request,
496        options: &Option<RequestOptions>,
497    ) -> Result<Response, ApiError> {
498        let max_retries = options
499            .as_ref()
500            .and_then(|opts| opts.max_retries)
501            .unwrap_or(self.config.max_retries);
502
503        let mut last_error = None;
504
505        for attempt in 0..=max_retries {
506            let cloned_request = request.try_clone().ok_or(ApiError::RequestClone)?;
507
508            match self.client.execute(cloned_request).await {
509                Ok(response) if response.status().is_success() => return Ok(response),
510                Ok(response)
511                    if attempt < max_retries
512                        && Self::is_retryable_status(response.status().as_u16()) =>
513                {
514                    // Exponential backoff for retryable HTTP status codes
515                    let delay = std::time::Duration::from_millis(100 * 2_u64.pow(attempt));
516                    tokio::time::sleep(delay).await;
517                }
518                Ok(response) => {
519                    let status_code = response.status().as_u16();
520                    let body = response.text().await.ok();
521                    return Err(ApiError::from_response(status_code, body.as_deref()));
522                }
523                Err(e) if attempt < max_retries => {
524                    last_error = Some(e);
525                    // Exponential backoff
526                    let delay = std::time::Duration::from_millis(100 * 2_u64.pow(attempt));
527                    tokio::time::sleep(delay).await;
528                }
529                Err(e) => return Err(ApiError::Network(e)),
530            }
531        }
532
533        Err(ApiError::Network(last_error.unwrap()))
534    }
535
536    fn is_retryable_status(status_code: u16) -> bool {
537        [408, 429].contains(&status_code) || status_code >= 500
538    }
539
540    async fn parse_response<T>(&self, response: Response) -> Result<T, ApiError>
541    where
542        T: DeserializeOwned,
543    {
544        let status = response.status().as_u16();
545        let text = response.text().await.map_err(ApiError::Network)?;
546
547        if text.is_empty() {
548            if status >= 400 {
549                return Err(ApiError::Http {
550                    status,
551                    message: String::new(),
552                });
553            }
554            return serde_json::from_value(serde_json::Value::Null).map_err(|_| ApiError::Http {
555                status,
556                message: String::new(),
557            });
558        }
559
560        serde_json::from_str(&text).map_err(ApiError::Serialization)
561    }
562
563    async fn parse_response_raw<T>(&self, response: Response) -> Result<RawResponse<T>, ApiError>
564    where
565        T: DeserializeOwned,
566    {
567        let status_code = response.status().as_u16();
568        let headers = response.headers().clone();
569        let text = response.text().await.map_err(ApiError::Network)?;
570
571        if text.is_empty() {
572            if status_code >= 400 {
573                return Err(ApiError::Http {
574                    status: status_code,
575                    message: String::new(),
576                });
577            }
578            return serde_json::from_value(serde_json::Value::Null)
579                .map(|body| RawResponse {
580                    body,
581                    status_code,
582                    headers,
583                })
584                .map_err(|_| ApiError::Http {
585                    status: status_code,
586                    message: String::new(),
587                });
588        }
589
590        let body: T = serde_json::from_str(&text).map_err(ApiError::Serialization)?;
591        Ok(RawResponse {
592            body,
593            status_code,
594            headers,
595        })
596    }
597
598    /// Execute a request and return a streaming response (for large file downloads)
599    ///
600    /// This method returns a `ByteStream` that can be used to download large files
601    /// efficiently without loading the entire content into memory. The stream can be
602    /// consumed chunk by chunk, written directly to disk, or collected into bytes.
603    ///
604    /// # Examples
605    ///
606    /// **Option 1: Collect all bytes into memory**
607    /// ```no_run
608    /// let stream = client.execute_stream_request(
609    ///     Method::GET,
610    ///     "/file",
611    ///     None,
612    ///     None,
613    ///     None,
614    /// ).await?;
615    ///
616    /// let bytes = stream.collect().await?;
617    /// ```
618    ///
619    /// **Option 2: Process chunks with try_next()**
620    /// ```no_run
621    /// let mut stream = client.execute_stream_request(
622    ///     Method::GET,
623    ///     "/large-file",
624    ///     None,
625    ///     None,
626    ///     None,
627    /// ).await?;
628    ///
629    /// while let Some(chunk) = stream.try_next().await? {
630    ///     process_chunk(&chunk);
631    /// }
632    /// ```
633    ///
634    /// **Option 3: Stream with futures::Stream trait**
635    /// ```no_run
636    /// use futures::StreamExt;
637    ///
638    /// let stream = client.execute_stream_request(
639    ///     Method::GET,
640    ///     "/large-file",
641    ///     None,
642    ///     None,
643    ///     None,
644    /// ).await?;
645    ///
646    /// let mut file = tokio::fs::File::create("output.mp4").await?;
647    /// let mut stream = std::pin::pin!(stream);
648    /// while let Some(chunk) = stream.next().await {
649    ///     let chunk = chunk?;
650    ///     tokio::io::AsyncWriteExt::write_all(&mut file, &chunk).await?;
651    /// }
652    /// ```
653    pub async fn execute_stream_request(
654        &self,
655        method: Method,
656        path: &str,
657        body: Option<serde_json::Value>,
658        query_params: Option<Vec<(String, String)>>,
659        options: Option<RequestOptions>,
660    ) -> Result<ByteStream, ApiError> {
661        let url = join_url(&self.config.base_url, path);
662        let mut request = self.client.request(method, &url);
663
664        // Apply query parameters if provided
665        if let Some(params) = query_params {
666            request = request.query(&params);
667        }
668
669        // Apply additional query parameters from options
670        if let Some(opts) = &options {
671            if !opts.additional_query_params.is_empty() {
672                request = request.query(&opts.additional_query_params);
673            }
674        }
675
676        // Apply body if provided
677        if let Some(body) = body {
678            request = request.json(&body);
679        }
680
681        // Build the request
682        let req = request.build().map_err(|e| ApiError::Network(e))?;
683
684        let response = self.send_request(req, &options).await?;
685
686        // Return streaming response
687        Ok(ByteStream::new(response))
688    }
689
690    /// Execute a streaming request with an explicit base URL override.
691    pub async fn execute_stream_request_with_base_url(
692        &self,
693        base_url: &str,
694        method: Method,
695        path: &str,
696        body: Option<serde_json::Value>,
697        query_params: Option<Vec<(String, String)>>,
698        options: Option<RequestOptions>,
699    ) -> Result<ByteStream, ApiError> {
700        let url = join_url(base_url, path);
701        let mut request = self.client.request(method, &url);
702
703        if let Some(params) = query_params {
704            request = request.query(&params);
705        }
706
707        if let Some(opts) = &options {
708            if !opts.additional_query_params.is_empty() {
709                request = request.query(&opts.additional_query_params);
710            }
711        }
712
713        if let Some(body) = body {
714            request = request.json(&body);
715        }
716
717        let req = request.build().map_err(|e| ApiError::Network(e))?;
718
719        let response = self.send_request(req, &options).await?;
720
721        Ok(ByteStream::new(response))
722    }
723}
724
725#[cfg(test)]
726mod tests {
727    use super::*;
728
729    #[test]
730    fn test_is_retryable_status() {
731        // Retryable 4xx
732        assert!(HttpClient::is_retryable_status(408));
733        assert!(HttpClient::is_retryable_status(429));
734
735        // Retryable 5xx (>= 500)
736        assert!(HttpClient::is_retryable_status(500));
737        assert!(HttpClient::is_retryable_status(501));
738        assert!(HttpClient::is_retryable_status(502));
739        assert!(HttpClient::is_retryable_status(503));
740        assert!(HttpClient::is_retryable_status(504));
741        assert!(HttpClient::is_retryable_status(599));
742
743        // Success and other 4xx codes are NOT retryable
744        assert!(!HttpClient::is_retryable_status(200));
745        assert!(!HttpClient::is_retryable_status(400));
746        assert!(!HttpClient::is_retryable_status(401));
747        assert!(!HttpClient::is_retryable_status(404));
748    }
749}
tracedb_sdk/core/http_client.rs

tracedb_sdk/core/
http_client.rs
