1#[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
21pub struct AsyncTesseraClient {
26 config: ClientConfig,
27 http: reqwest::Client,
28}
29
30impl AsyncTesseraClient {
31 pub fn new(api_key: Option<&str>) -> Result<Self, TesseraError> {
38 Self::from_config(ClientConfig::new(api_key)?)
39 }
40
41 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 #[must_use]
58 pub fn config(&self) -> &ClientConfig {
59 &self.config
60 }
61
62 pub fn close(&mut self) {
67 self.http = reqwest::Client::new();
68 }
69
70 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 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 pub async fn datasets(&self) -> Result<DatasetsResponse, TesseraError> {
95 let response = self.request(&datasets_request()).await?;
96 self.json(response).await
97 }
98
99 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 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 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 #[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 #[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 #[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 #[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
233async 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
282async 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}