Skip to main content

tessera/
async_client.rs

1//! Asyncio-native Tessera API client (`async fn` surface).
2
3#[cfg(feature = "polars")]
4use polars::prelude::{DataFrame, LazyFrame};
5
6use crate::base::{
7    PreparedRequest, datasets_request, download_request, parse_retry_after, partitions_request,
8    should_retry,
9};
10use crate::config::{ClientConfig, backoff_delay};
11use crate::error::{TesseraError, error_from_response};
12use crate::models::{
13    DatasetsResponse, DownloadResponse, IntoCoins, IntoMonths, PartitionRef, PartitionsResponse,
14};
15#[cfg(any(feature = "polars", feature = "duckdb"))]
16use crate::readers::ResolvedPartition;
17use crate::readers::expand_refs;
18#[cfg(any(feature = "polars", feature = "duckdb"))]
19use crate::resolver::resolve_async;
20
21/// An async client for the Tessera API.
22///
23/// Mirror of [`crate::TesseraClient`] with `async fn`; use it inside tokio
24/// runtimes (including when the sync client would panic).
25pub struct AsyncTesseraClient {
26    config: ClientConfig,
27    http: reqwest::Client,
28}
29
30impl AsyncTesseraClient {
31    /// Create a client with defaults, resolving the key from `api_key` or
32    /// `$TESSERA_API_KEY`.
33    ///
34    /// # Errors
35    ///
36    /// Returns [`TesseraError::Configuration`] when no API key can be resolved.
37    pub fn new(api_key: Option<&str>) -> Result<Self, TesseraError> {
38        Self::from_config(ClientConfig::new(api_key)?)
39    }
40
41    /// Create a client from a fully-specified [`ClientConfig`].
42    ///
43    /// # Errors
44    ///
45    /// Returns [`TesseraError::Network`] when the HTTP client cannot be built.
46    pub fn from_config(config: ClientConfig) -> Result<Self, TesseraError> {
47        let http = reqwest::Client::builder()
48            .redirect(reqwest::redirect::Policy::none())
49            .default_headers(config.auth_headers())
50            .timeout(config.timeout)
51            .build()
52            .map_err(|err| TesseraError::Network(err.to_string()))?;
53        Ok(Self { config, http })
54    }
55
56    /// The resolved configuration this client was built with.
57    #[must_use]
58    pub fn config(&self) -> &ClientConfig {
59        &self.config
60    }
61
62    /// Close the client.
63    ///
64    /// `reqwest::Client` tears its pool down on drop; taking the client out
65    /// eagerly releases the handle for parity with the Python `aclose()`.
66    pub fn close(&mut self) {
67        self.http = reqwest::Client::new();
68    }
69
70    /// Send a prepared request, retrying transient failures.
71    async fn request(&self, prepared: &PreparedRequest) -> Result<reqwest::Response, TesseraError> {
72        send_with_retries(
73            &self.http,
74            &self.config.base_url,
75            self.config.max_retries,
76            prepared,
77        )
78        .await
79    }
80
81    /// Parse a success response body as JSON.
82    async fn json<T: serde::de::DeserializeOwned>(
83        &self,
84        response: reqwest::Response,
85    ) -> Result<T, TesseraError> {
86        parse_json(response).await
87    }
88
89    /// List every dataset visible to your plan.
90    ///
91    /// # Errors
92    ///
93    /// Returns any [`TesseraError`] the API or transport raises.
94    pub async fn datasets(&self) -> Result<DatasetsResponse, TesseraError> {
95        let response = self.request(&datasets_request()).await?;
96        self.json(response).await
97    }
98
99    /// List the partitions of `asset`, optionally filtered by coin/month.
100    ///
101    /// # Errors
102    ///
103    /// Returns any [`TesseraError`] the API or transport raises.
104    pub async fn partitions(
105        &self,
106        asset: &str,
107        coin: Option<&str>,
108        month: Option<&str>,
109    ) -> Result<PartitionsResponse, TesseraError> {
110        let response = self
111            .request(&partitions_request(asset, coin, month))
112            .await?;
113        self.json(response).await
114    }
115
116    /// Mint a short-lived presigned download URL for one partition.
117    ///
118    /// # Errors
119    ///
120    /// Returns any [`TesseraError`] the API or transport raises.
121    pub async fn download_url(
122        &self,
123        asset: &str,
124        coin: &str,
125        month: &str,
126    ) -> Result<DownloadResponse, TesseraError> {
127        let response = self.request(&download_request(asset, coin, month)).await?;
128        self.json(response).await
129    }
130
131    /// Expand `(asset, coins, months)` into concrete partition references.
132    ///
133    /// # Errors
134    ///
135    /// Returns [`TesseraError::InvalidArgument`] on empty/invalid arguments.
136    pub fn partition_refs(
137        &self,
138        asset: &str,
139        coin: impl IntoCoins,
140        month: impl IntoMonths,
141    ) -> Result<Vec<PartitionRef>, TesseraError> {
142        let _ = self;
143        expand_refs(asset, coin, month)
144    }
145
146    /// Resolve presigned URLs for every ref, concurrently and order-preserving.
147    #[cfg(any(feature = "polars", feature = "duckdb"))]
148    async fn resolve(&self, refs: &[PartitionRef]) -> Result<Vec<ResolvedPartition>, TesseraError> {
149        let http = self.http.clone();
150        let base_url = self.config.base_url.clone();
151        let max_retries = self.config.max_retries;
152        resolve_async(
153            move |partition| {
154                let http = http.clone();
155                let base_url = base_url.clone();
156                async move {
157                    let prepared =
158                        download_request(&partition.asset, &partition.coin, &partition.month);
159                    let response =
160                        send_with_retries(&http, &base_url, max_retries, &prepared).await?;
161                    parse_json(response).await.map(|d: DownloadResponse| d.url)
162                }
163            },
164            refs,
165        )
166        .await
167    }
168
169    /// Lazily scan one or more partitions into a Polars `LazyFrame`.
170    ///
171    /// # Errors
172    ///
173    /// Returns any [`TesseraError`] the API or transport raises.
174    #[cfg(feature = "polars")]
175    pub async fn scan(
176        &self,
177        asset: &str,
178        coin: impl IntoCoins,
179        month: impl IntoMonths,
180        columns: Option<&[&str]>,
181    ) -> Result<LazyFrame, TesseraError> {
182        let parts = self.resolve(&expand_refs(asset, coin, month)?).await?;
183        crate::readers::polars::build_lazyframe(&parts, columns)
184    }
185
186    /// Eagerly read one or more partitions into a Polars `DataFrame`.
187    ///
188    /// # Errors
189    ///
190    /// Returns any [`TesseraError`] the API or transport raises, including
191    /// [`TesseraError::PresignExpired`] for rejected presigned URLs.
192    #[cfg(feature = "polars")]
193    pub async fn read(
194        &self,
195        asset: &str,
196        coin: impl IntoCoins,
197        month: impl IntoMonths,
198        columns: Option<&[&str]>,
199    ) -> Result<DataFrame, TesseraError> {
200        let parts = self.resolve(&expand_refs(asset, coin, month)?).await?;
201        let lazy = crate::readers::polars::build_lazyframe(&parts, columns)?;
202        crate::readers::polars::collect(lazy)
203    }
204
205    /// Open one or more partitions as an in-memory DuckDB connection exposing
206    /// a `tessera` view for SQL querying.
207    ///
208    /// # Errors
209    ///
210    /// Returns any [`TesseraError`] the API or transport raises.
211    #[cfg(feature = "duckdb")]
212    pub async fn to_duckdb(
213        &self,
214        asset: &str,
215        coin: impl IntoCoins,
216        month: impl IntoMonths,
217        columns: Option<&[&str]>,
218    ) -> Result<duckdb::Connection, TesseraError> {
219        let parts = self.resolve(&expand_refs(asset, coin, month)?).await?;
220        let owned_columns: Option<Vec<String>> =
221            columns.map(|cols| cols.iter().map(|c| (*c).to_string()).collect());
222        tokio::task::spawn_blocking(move || {
223            let borrowed: Option<Vec<&str>> = owned_columns
224                .as_deref()
225                .map(|cols| cols.iter().map(String::as_str).collect());
226            crate::readers::duckdb::build_relation(&parts, borrowed.as_deref())
227        })
228        .await
229        .map_err(|err| TesseraError::Network(format!("duckdb task failed: {err}")))?
230    }
231}
232
233/// Send a request with the shared retry policy (transient network errors and
234/// retryable statuses back off exponentially, honouring `Retry-After`).
235async fn send_with_retries(
236    http: &reqwest::Client,
237    base_url: &str,
238    max_retries: u32,
239    prepared: &PreparedRequest,
240) -> Result<reqwest::Response, TesseraError> {
241    for attempt in 0..=max_retries {
242        let last = attempt == max_retries;
243        let sent = http
244            .get(format!("{base_url}{}", prepared.path))
245            .query(&prepared.params)
246            .send()
247            .await;
248        match sent {
249            Err(err) if err.is_connect() || err.is_timeout() || err.is_request() => {
250                if last {
251                    return Err(TesseraError::Network(format!(
252                        "network error contacting Tessera: {err}"
253                    )));
254                }
255                tokio::time::sleep(backoff_delay(attempt, None)).await;
256            }
257            Ok(response) if !last && should_retry(response.status().as_u16()) => {
258                let retry_after = parse_retry_after(response.headers());
259                tokio::time::sleep(backoff_delay(attempt, retry_after)).await;
260            }
261            Ok(response) => {
262                if response.status().is_success() {
263                    return Ok(response);
264                }
265                let status = response.status().as_u16();
266                let body = response
267                    .bytes()
268                    .await
269                    .map_err(|err| TesseraError::Network(err.to_string()))?;
270                return Err(error_from_response(status, &body));
271            }
272            Err(err) => {
273                return Err(TesseraError::Network(format!(
274                    "network error contacting Tessera: {err}"
275                )));
276            }
277        }
278    }
279    unreachable!("retry loop always returns")
280}
281
282/// Parse a success response body as JSON.
283async fn parse_json<T: serde::de::DeserializeOwned>(
284    response: reqwest::Response,
285) -> Result<T, TesseraError> {
286    let body = response
287        .bytes()
288        .await
289        .map_err(|err| TesseraError::Network(err.to_string()))?;
290    serde_json::from_slice(&body).map_err(|err| TesseraError::Network(err.to_string()))
291}
292
293impl std::fmt::Debug for AsyncTesseraClient {
294    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295        f.debug_struct("AsyncTesseraClient")
296            .field("base_url", &self.config.base_url)
297            .field("timeout", &self.config.timeout)
298            .finish_non_exhaustive()
299    }
300}