influxdb3_client/
client.rs

1use std::collections::HashMap;
2use std::future::IntoFuture;
3use std::io::Write as IoWrite;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7
8use bytes::Bytes;
9use futures_util::stream::{self, StreamExt, TryStreamExt};
10use reqwest::{
11    header::{AUTHORIZATION, CONTENT_ENCODING, CONTENT_TYPE, RETRY_AFTER},
12    Client as HttpClient, ClientBuilder, Response,
13};
14use tokio::sync::OnceCell;
15
16use crate::{
17    config::ClientConfig,
18    error::{Error, LineError, PartialWriteError},
19    flight::{BatchStream, FlightQueryClient},
20    precision::Precision,
21    query::{QueryOptions, QueryParameters, QueryResult, QueryType},
22    retry::{self, RetryConfig},
23    write::{WriteInput, WriteOptions},
24    Result,
25};
26
27/// Async client for InfluxDB 3 Core and Enterprise.
28///
29/// See the crate-level docs for end-to-end examples.
30pub struct Client {
31    config: ClientConfig,
32    http: HttpClient,
33    /// Lazy: connected on first query.  `OnceCell` retries on init failure.
34    flight: OnceCell<FlightQueryClient>,
35}
36
37impl Client {
38    /// Create a client from a [`ClientConfig`].
39    ///
40    /// The Arrow Flight gRPC channel is opened lazily on the first query call,
41    /// so this constructor never fails due to gRPC connectivity.
42    pub async fn new(config: ClientConfig) -> Result<Self> {
43        let http = build_http_client(&config)?;
44        Ok(Client {
45            config,
46            http,
47            flight: OnceCell::new(),
48        })
49    }
50
51    /// Parse a connection string and create a client.
52    pub async fn from_connection_string(cs: &str) -> Result<Self> {
53        Client::new(ClientConfig::from_connection_string(cs)?).await
54    }
55
56    /// Read configuration from environment variables and create a client.
57    ///
58    /// See [`ClientConfig::from_env`] for the supported variables.
59    pub async fn from_env() -> Result<Self> {
60        Client::new(ClientConfig::from_env()?).await
61    }
62
63    /// Return a reference to the underlying config.
64    pub fn config(&self) -> &ClientConfig {
65        &self.config
66    }
67
68    /// Start a write request.
69    ///
70    /// `data` may be any [`WriteInput`]: a `&str` / `String` of pre-formatted
71    /// line protocol, a `Vec<Point>` / `&[Point]`, a [`DataFrameWrite`] (polars
72    /// feature), or your own type that implements the trait.
73    ///
74    /// Returns a [`WriteRequest`] builder; chain options, then `.await`.
75    /// See the crate-level docs for examples.
76    ///
77    /// [`DataFrameWrite`]: crate::write_dataframe::DataFrameWrite
78    pub fn write<W: WriteInput>(&self, data: W) -> WriteRequest<'_, W> {
79        WriteRequest {
80            client: self,
81            data: Some(data),
82            options: self.config.write_options.clone(),
83            retry: None,
84        }
85    }
86
87    /// Start a SQL query.  Sugar for `query(q, QueryType::Sql)`.
88    pub fn sql(&self, q: impl Into<String>) -> QueryRequest<'_> {
89        self.query(q, QueryType::Sql)
90    }
91
92    /// Start an InfluxQL query.  Sugar for `query(q, QueryType::InfluxQL)`.
93    pub fn influxql(&self, q: impl Into<String>) -> QueryRequest<'_> {
94        self.query(q, QueryType::InfluxQL)
95    }
96
97    /// Start a query.  Returns a [`QueryRequest`] builder.
98    pub fn query(&self, q: impl Into<String>, language: QueryType) -> QueryRequest<'_> {
99        QueryRequest {
100            client: self,
101            query: q.into(),
102            query_type: language,
103            params: QueryParameters::new(),
104            headers: HashMap::new(),
105            retry: None,
106        }
107    }
108
109    /// Ping the server and return its version string.
110    pub async fn ping(&self) -> Result<String> {
111        let url = format!("{}/ping", self.config.host_url());
112        let mut req = self.http.get(&url);
113        if let Some(auth) = self.config.authorization_header()? {
114            req = req.header(AUTHORIZATION, auth);
115        }
116        let resp = req.send().await?;
117        let version = resp
118            .headers()
119            .get("x-influxdb-version")
120            .and_then(|v| v.to_str().ok())
121            .unwrap_or("unknown")
122            .to_owned();
123        Ok(version)
124    }
125
126    /// Internal: resolve the Flight client (lazy-init on first call).
127    async fn flight(&self) -> Result<&FlightQueryClient> {
128        self.flight
129            .get_or_try_init(|| async {
130                FlightQueryClient::new(
131                    &self.config.host,
132                    self.config.token.as_deref(),
133                    &self.config.auth_scheme,
134                    self.config.ssl_roots_path.as_deref(),
135                    self.config.query_timeout,
136                )
137                .await
138            })
139            .await
140    }
141
142    /// Internal: send one LP batch, retrying transient failures per `policy`.
143    async fn send_lp(
144        &self,
145        body: Vec<u8>,
146        opts: &WriteOptions,
147        policy: &RetryConfig,
148    ) -> Result<()> {
149        let db = &self.config.database;
150
151        let (url, mut params) = if opts.use_v2_api {
152            let url = format!("{}/api/v2/write", self.config.host_url());
153            let mut p = vec![("bucket", db.clone())];
154            if let Some(org) = &self.config.org {
155                p.push(("org", org.clone()));
156            }
157            (url, p)
158        } else {
159            let url = format!("{}/api/v3/write_lp", self.config.host_url());
160            let mut p = vec![("db", db.clone())];
161            p.push(("accept_partial", opts.accept_partial.to_string()));
162            if opts.no_sync {
163                p.push(("no_sync", "true".to_string()));
164            }
165            (url, p)
166        };
167
168        params.push(("precision", opts.precision.as_str().to_string()));
169
170        // Compress once; each attempt re-sends the same (Arc-backed) Bytes.
171        let (final_body, compressed) = maybe_gzip(body, opts.gzip_threshold).await?;
172
173        let mut base = self
174            .http
175            .post(&url)
176            .query(&params)
177            .header(CONTENT_TYPE, "text/plain; charset=utf-8");
178        if compressed {
179            base = base.header(CONTENT_ENCODING, "gzip");
180        }
181        if let Some(auth) = self.config.authorization_header()? {
182            base = base.header(AUTHORIZATION, auth);
183        }
184        for (k, v) in &self.config.headers {
185            base = base.header(k, v);
186        }
187        let base = base.body(final_body);
188
189        let start = Instant::now();
190        let mut attempt: u32 = 0;
191        loop {
192            let req = base.try_clone().expect("write body is cloneable");
193            match send_write_once(req).await {
194                WriteAttempt::Ok => return Ok(()),
195                WriteAttempt::Fatal(e) => return Err(e),
196                WriteAttempt::Retry { after, last } => {
197                    if attempt >= policy.max_retries {
198                        return Err(last);
199                    }
200                    let delay = after
201                        .filter(|_| policy.honor_retry_after)
202                        .unwrap_or_else(|| policy.backoff(attempt));
203                    if let Some(budget) = policy.max_elapsed {
204                        if start.elapsed() + delay > budget {
205                            return Err(last);
206                        }
207                    }
208                    tokio::time::sleep(delay).await;
209                    attempt += 1;
210                }
211            }
212        }
213    }
214}
215
216/// Builder produced by [`Client::write`]; chain options, then `.await`.
217pub struct WriteRequest<'a, W: WriteInput> {
218    client: &'a Client,
219    data: Option<W>,
220    options: WriteOptions,
221    retry: Option<RetryConfig>,
222}
223
224impl<'a, W: WriteInput> WriteRequest<'a, W> {
225    pub fn precision(mut self, p: Precision) -> Self {
226        self.options.precision = p;
227        self
228    }
229    pub fn no_sync(mut self) -> Self {
230        self.options.no_sync = true;
231        self
232    }
233    pub fn accept_partial(mut self, accept: bool) -> Self {
234        self.options.accept_partial = accept;
235        self
236    }
237    pub fn use_v2_api(mut self) -> Self {
238        self.options.use_v2_api = true;
239        self
240    }
241    pub fn batch_size(mut self, n: usize) -> Self {
242        self.options.batch_size = n;
243        self
244    }
245    pub fn max_inflight(mut self, n: usize) -> Self {
246        self.options.max_inflight = n;
247        self
248    }
249    pub fn default_tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
250        self.options.default_tags.insert(key.into(), value.into());
251        self
252    }
253    /// Compress bodies larger than `t` bytes; `None` disables compression.
254    /// Disable or raise this for high-throughput ingest over a fast LAN where
255    /// gzip CPU outweighs the bandwidth saved. See [`WriteOptions::gzip_threshold`].
256    pub fn gzip_threshold(mut self, t: Option<usize>) -> Self {
257        self.options.gzip_threshold = t;
258        self
259    }
260    pub fn tag_order(mut self, order: impl IntoIterator<Item = impl Into<String>>) -> Self {
261        self.options.tag_order = order.into_iter().map(Into::into).collect();
262        self
263    }
264
265    /// Replace the underlying options wholesale.
266    pub fn with_options(mut self, opts: WriteOptions) -> Self {
267        self.options = opts;
268        self
269    }
270
271    /// Override the retry policy for this write (defaults to the client's).
272    pub fn retry(mut self, policy: RetryConfig) -> Self {
273        self.retry = Some(policy);
274        self
275    }
276
277    /// Disable retries for this write.
278    pub fn no_retry(mut self) -> Self {
279        self.retry = Some(RetryConfig::disabled());
280        self
281    }
282}
283
284impl<'a, W: WriteInput + Send + 'a> IntoFuture for WriteRequest<'a, W> {
285    type Output = Result<()>;
286    type IntoFuture = Pin<Box<dyn std::future::Future<Output = Self::Output> + Send + 'a>>;
287
288    fn into_future(mut self) -> Self::IntoFuture {
289        let client = self.client;
290        let data = self.data.take().expect("data already taken");
291        let options = self.options;
292        let policy = self.retry.unwrap_or_else(|| client.config.retry.clone());
293        Box::pin(async move {
294            options.validate()?;
295            let max_inflight = options.max_inflight.max(1);
296            let batches = data.into_lp_batches(&options);
297
298            if max_inflight == 1 {
299                for batch in batches {
300                    let bytes = batch?;
301                    client.send_lp(bytes, &options, &policy).await?;
302                }
303                return Ok(());
304            }
305
306            let options = Arc::new(options);
307            let policy = Arc::new(policy);
308            stream::iter(batches)
309                .map(|b| b.map(|bytes| (bytes, Arc::clone(&options), Arc::clone(&policy))))
310                .try_for_each_concurrent(Some(max_inflight), |(bytes, opts, pol)| async move {
311                    client.send_lp(bytes, &opts, &pol).await
312                })
313                .await
314        })
315    }
316}
317
318/// Builder produced by [`Client::sql`], [`Client::influxql`], or
319/// [`Client::query`]; chain options, then `.await` (for a collected
320/// [`QueryResult`]) or `.stream()` (for a streaming [`BatchStream`]).
321pub struct QueryRequest<'a> {
322    client: &'a Client,
323    query: String,
324    query_type: QueryType,
325    params: QueryParameters,
326    headers: HashMap<String, String>,
327    retry: Option<RetryConfig>,
328}
329
330impl<'a> QueryRequest<'a> {
331    /// Add a single named parameter.
332    pub fn param(mut self, key: impl Into<String>, value: impl Into<serde_json::Value>) -> Self {
333        self.params.insert(key.into(), value.into());
334        self
335    }
336
337    /// Add multiple named parameters from an iterable.
338    pub fn params<K, V, I>(mut self, params: I) -> Self
339    where
340        I: IntoIterator<Item = (K, V)>,
341        K: Into<String>,
342        V: Into<serde_json::Value>,
343    {
344        for (k, v) in params {
345            self.params.insert(k.into(), v.into());
346        }
347        self
348    }
349
350    /// Add a gRPC metadata header sent with the Flight DoGet request.
351    pub fn header(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
352        self.headers.insert(k.into(), v.into());
353        self
354    }
355
356    /// Override the retry policy for this query (defaults to the client's).
357    pub fn retry(mut self, policy: RetryConfig) -> Self {
358        self.retry = Some(policy);
359        self
360    }
361
362    /// Disable retries for this query.
363    pub fn no_retry(mut self) -> Self {
364        self.retry = Some(RetryConfig::disabled());
365        self
366    }
367
368    /// Open the query as a streaming [`BatchStream`] instead of collecting.
369    /// Use this for results too large to materialise in memory.
370    ///
371    /// Only the connection setup is retried; once batches start arriving a
372    /// stream can't be replayed, so mid-stream errors propagate to the caller.
373    pub async fn stream(self) -> Result<BatchStream> {
374        let policy = self
375            .retry
376            .clone()
377            .unwrap_or_else(|| self.client.config.retry.clone());
378        let opts = QueryOptions {
379            query_type: self.query_type,
380            headers: self.headers,
381        };
382        let params = (!self.params.is_empty()).then_some(self.params);
383
384        let mut attempt: u32 = 0;
385        loop {
386            let result = async {
387                let flight = self.client.flight().await?;
388                flight
389                    .stream(
390                        &self.query,
391                        &self.client.config.database,
392                        &opts,
393                        params.as_ref(),
394                    )
395                    .await
396            }
397            .await;
398
399            match result {
400                Ok(s) => return Ok(s),
401                Err(e) => {
402                    if attempt >= policy.max_retries || !is_retryable_query_err(&e) {
403                        return Err(e);
404                    }
405                    tokio::time::sleep(policy.backoff(attempt)).await;
406                    attempt += 1;
407                }
408            }
409        }
410    }
411}
412
413/// Whether a query error is a transient failure worth retrying.
414fn is_retryable_query_err(e: &Error) -> bool {
415    match e {
416        Error::Flight(status) => retry::retryable_tonic(status.code()),
417        Error::Transport(_) => true,
418        Error::Timeout(_) => true,
419        _ => false,
420    }
421}
422
423impl<'a> IntoFuture for QueryRequest<'a> {
424    type Output = Result<QueryResult>;
425    type IntoFuture = Pin<Box<dyn std::future::Future<Output = Self::Output> + Send + 'a>>;
426
427    fn into_future(self) -> Self::IntoFuture {
428        Box::pin(async move {
429            let timeout = self.client.config.query_timeout;
430            let policy = self
431                .retry
432                .clone()
433                .unwrap_or_else(|| self.client.config.retry.clone());
434            let opts = QueryOptions {
435                query_type: self.query_type,
436                headers: self.headers,
437            };
438            let params = (!self.params.is_empty()).then_some(self.params);
439
440            let start = Instant::now();
441            let mut attempt: u32 = 0;
442            loop {
443                // `timeout` is per attempt; collected queries are read-only and
444                // safe to re-issue in full.
445                let fut = async {
446                    let flight = self.client.flight().await?;
447                    flight
448                        .query(
449                            &self.query,
450                            &self.client.config.database,
451                            &opts,
452                            params.as_ref(),
453                        )
454                        .await
455                };
456                let outcome = match tokio::time::timeout(timeout, fut).await {
457                    Ok(inner) => inner,
458                    Err(_) => Err(Error::Timeout(timeout)),
459                };
460
461                match outcome {
462                    Ok(result) => return Ok(result),
463                    Err(e) => {
464                        if attempt >= policy.max_retries || !is_retryable_query_err(&e) {
465                            return Err(e);
466                        }
467                        let delay = policy.backoff(attempt);
468                        if let Some(budget) = policy.max_elapsed {
469                            if start.elapsed() + delay > budget {
470                                return Err(e);
471                            }
472                        }
473                        tokio::time::sleep(delay).await;
474                        attempt += 1;
475                    }
476                }
477            }
478        })
479    }
480}
481
482fn build_http_client(config: &ClientConfig) -> Result<HttpClient> {
483    let mut builder = ClientBuilder::new()
484        .timeout(config.write_timeout)
485        .pool_idle_timeout(config.idle_connection_timeout)
486        .pool_max_idle_per_host(config.max_idle_connections)
487        .gzip(true)
488        .use_rustls_tls();
489
490    if let Some(proxy_url) = &config.proxy {
491        let proxy = reqwest::Proxy::all(proxy_url)
492            .map_err(|e| Error::Config(format!("invalid proxy URL: {e}")))?;
493        builder = builder.proxy(proxy);
494    }
495
496    if let Some(roots_path) = &config.ssl_roots_path {
497        let pem = std::fs::read(roots_path)
498            .map_err(|e| Error::Config(format!("cannot read SSL roots '{roots_path}': {e}")))?;
499        let cert = reqwest::tls::Certificate::from_pem(&pem)
500            .map_err(|e| Error::Config(format!("invalid SSL roots PEM: {e}")))?;
501        builder = builder.add_root_certificate(cert);
502    }
503
504    Ok(builder.build()?)
505}
506
507/// Threshold above which gzip compression runs on a blocking thread pool.
508/// For smaller bodies the spawn_blocking overhead dominates the compression cost.
509const SPAWN_BLOCKING_GZIP_THRESHOLD: usize = 64 * 1024;
510
511/// Maybe gzip-compress a body.  Returns `(body_bytes, was_compressed)`.
512async fn maybe_gzip(data: Vec<u8>, threshold: Option<usize>) -> Result<(Bytes, bool)> {
513    let should_compress = matches!(threshold, Some(t) if data.len() > t);
514    if !should_compress {
515        return Ok((Bytes::from(data), false));
516    }
517
518    if data.len() < SPAWN_BLOCKING_GZIP_THRESHOLD {
519        let compressed = gzip_compress(data)?;
520        return Ok((Bytes::from(compressed), true));
521    }
522
523    let compressed = tokio::task::spawn_blocking(move || gzip_compress(data))
524        .await
525        .map_err(|e| Error::Config(format!("gzip task join error: {e}")))??;
526    Ok((Bytes::from(compressed), true))
527}
528
529fn gzip_compress(data: Vec<u8>) -> Result<Vec<u8>> {
530    let mut encoder = flate2::write::GzEncoder::new(
531        Vec::with_capacity(data.len() / 2),
532        flate2::Compression::default(),
533    );
534    encoder
535        .write_all(&data)
536        .map_err(|e| Error::Config(format!("gzip encoding failed: {e}")))?;
537    encoder
538        .finish()
539        .map_err(|e| Error::Config(format!("gzip finalization failed: {e}")))
540}
541
542/// Outcome of a single write attempt, before any backoff decision.
543enum WriteAttempt {
544    Ok,
545    /// Transient, safe to retry. `last` is surfaced if retries are exhausted;
546    /// `after` is the server-suggested delay from `Retry-After`, if any.
547    Retry {
548        after: Option<Duration>,
549        last: Error,
550    },
551    /// Deterministic; do not retry (auth, bad request, partial write, etc.).
552    Fatal(Error),
553}
554
555/// Send one write request and classify the result into a [`WriteAttempt`].
556async fn send_write_once(req: reqwest::RequestBuilder) -> WriteAttempt {
557    match req.send().await {
558        Err(e) => {
559            let retryable = retry::retryable_reqwest(&e);
560            let err = Error::Http(e);
561            if retryable {
562                WriteAttempt::Retry {
563                    after: None,
564                    last: err,
565                }
566            } else {
567                WriteAttempt::Fatal(err)
568            }
569        }
570        Ok(resp) => {
571            let status = resp.status();
572            if status.is_success() {
573                return WriteAttempt::Ok;
574            }
575            let code = status.as_u16();
576            let retry_after = resp
577                .headers()
578                .get(RETRY_AFTER)
579                .and_then(|v| v.to_str().ok())
580                .and_then(retry::parse_retry_after);
581            // A partial write is a 400, so it falls through to Fatal here:
582            // a deterministic data error, never a transient one.
583            let err = parse_write_error(code, resp).await;
584            if retry::retryable_status(code) {
585                WriteAttempt::Retry {
586                    after: retry_after,
587                    last: err,
588                }
589            } else {
590                WriteAttempt::Fatal(err)
591            }
592        }
593    }
594}
595
596/// Parse a non-2xx write response body into an [`Error`] (partial write vs server).
597async fn parse_write_error(code: u16, resp: Response) -> Error {
598    let body = resp.text().await.unwrap_or_default();
599
600    if let Ok(v) = serde_json::from_str::<serde_json::Value>(&body) {
601        let is_partial = v
602            .get("error")
603            .and_then(|e| e.as_str())
604            .map(|s| s.contains("partial write"))
605            .unwrap_or(false);
606
607        if is_partial && v.get("data").and_then(|d| d.as_array()).is_some() {
608            return Error::PartialWrite(PartialWriteError {
609                line_errors: parse_line_errors(&v),
610            });
611        }
612
613        let msg = v
614            .get("error")
615            .or_else(|| v.get("message"))
616            .and_then(|m| m.as_str())
617            .unwrap_or(&body)
618            .to_owned();
619
620        return Error::Server { code, message: msg };
621    }
622
623    Error::Server {
624        code,
625        message: body,
626    }
627}
628
629fn parse_line_errors(v: &serde_json::Value) -> Vec<LineError> {
630    v.get("data")
631        .and_then(|d| d.as_array())
632        .map(|arr| {
633            arr.iter()
634                .filter_map(|e| {
635                    Some(LineError {
636                        line: e.get("line_number")?.as_u64()?,
637                        message: e.get("error_message")?.as_str()?.to_owned(),
638                        original_line: e
639                            .get("original_line")
640                            .and_then(|s| s.as_str())
641                            .map(str::to_owned),
642                    })
643                })
644                .collect()
645        })
646        .unwrap_or_default()
647}
influxdb3_client/client.rs

influxdb3_client/
client.rs
