clia_influxdb2/api/
write.rs

1//! Write API
2
3use crate::models::WriteDataPoint;
4use crate::{Client, Http, RequestError, ReqwestProcessing};
5
6use bytes::BufMut;
7use futures::{Stream, StreamExt};
8use reqwest::header::HeaderMap;
9use reqwest::{Body, Method, StatusCode};
10use snafu::ResultExt;
11use std::io::{self, Write};
12
13impl Client {
14    /// Write line protocol data to the specified organization and bucket.
15    /// This method writes with default timestamp precision (nanoseconds).
16    /// Use write_line_protocol_with_precision if you want to write with a different precision.
17    pub async fn write_line_protocol(
18        &self,
19        org: &str,
20        bucket: &str,
21        body: impl Into<Body> + Send,
22    ) -> Result<(), RequestError> {
23        self.write_line_protocol_with_precision(org, bucket, body, TimestampPrecision::Nanoseconds)
24            .await
25    }
26
27    /// Write line protocol data to the specified organization and bucket.
28    pub async fn write_line_protocol_with_precision(
29        &self,
30        org: &str,
31        bucket: &str,
32        body: impl Into<Body> + Send,
33        precision: TimestampPrecision,
34    ) -> Result<(), RequestError> {
35        self.write_line_protocol_with_precision_headers(
36            org,
37            bucket,
38            body,
39            precision,
40            HeaderMap::new(),
41        )
42        .await
43    }
44
45    async fn write_line_protocol_with_precision_headers(
46        &self,
47        org: &str,
48        bucket: &str,
49        body: impl Into<Body> + Send,
50        precision: TimestampPrecision,
51        headers: HeaderMap,
52    ) -> Result<(), RequestError> {
53        let body = body.into();
54        let write_url = self.url("/api/v2/write");
55
56        let response = self
57            .request(Method::POST, &write_url)
58            .headers(headers)
59            .query(&[
60                ("bucket", bucket),
61                ("org", org),
62                ("precision", precision.api_short_name()),
63            ])
64            .body(body)
65            .send()
66            .await
67            .context(ReqwestProcessing)?;
68
69        if response.status() != StatusCode::NO_CONTENT {
70            let status = response.status();
71            let text = response.text().await.context(ReqwestProcessing)?;
72            Http { status, text }.fail()?;
73        }
74
75        Ok(())
76    }
77
78    /// Write a `Stream` of `DataPoint`s to the specified bucket.
79    ///
80    /// This method writes with default timestamp precision (nanoseconds).
81    /// Use write_with_precision if you want to write with a different precision.
82    pub async fn write(
83        &self,
84        bucket: &str,
85        body: impl Stream<Item = impl WriteDataPoint> + Send + Sync + 'static,
86    ) -> Result<(), RequestError> {
87        self.write_with_precision(bucket, body, TimestampPrecision::Nanoseconds)
88            .await
89    }
90
91    /// Write a `Stream` of `DataPoint`s to the specified organization and
92    /// bucket.
93    pub async fn write_with_precision(
94        &self,
95        bucket: &str,
96        body: impl Stream<Item = impl WriteDataPoint> + Send + Sync + 'static,
97        timestamp_precision: TimestampPrecision,
98    ) -> Result<(), RequestError> {
99        let mut buffer = bytes::BytesMut::new();
100
101        let body = body.map(move |point| {
102            let mut w = (&mut buffer).writer();
103            point.write_data_point_to(&mut w)?;
104            w.flush()?;
105            Ok::<_, io::Error>(buffer.split().freeze())
106        });
107
108        #[cfg(feature = "gzip")]
109        {
110            use crate::Compression;
111            use async_compression::tokio::bufread::GzipEncoder;
112            use async_compression::Level;
113            use reqwest::header::HeaderValue;
114            use tokio_util::io::{ReaderStream, StreamReader};
115
116            match self.compression {
117                Compression::Gzip => {
118                    let encoder = GzipEncoder::with_quality(StreamReader::new(body), Level::Best);
119                    let body: Body = Body::wrap_stream(ReaderStream::new(encoder));
120
121                    let mut headers = HeaderMap::new();
122                    headers.insert("Content-Encoding", HeaderValue::from_static("gzip"));
123
124                    return self
125                        .write_line_protocol_with_precision_headers(
126                            &self.org,
127                            bucket,
128                            body,
129                            timestamp_precision,
130                            headers,
131                        )
132                        .await;
133                }
134                Compression::None => {
135                    // fall through
136                }
137            }
138        }
139        let body: Body = Body::wrap_stream(body);
140
141        self.write_line_protocol_with_precision(&self.org, bucket, body, timestamp_precision)
142            .await
143    }
144}
145
146/// Possible timestamp precisions.
147#[derive(Debug, PartialEq, Copy, Clone)]
148pub enum TimestampPrecision {
149    /// Seconds timestamp precision
150    Seconds,
151    /// Milliseconds timestamp precision
152    Milliseconds,
153    /// Microseconds timestamp precision
154    Microseconds,
155    /// Nanoseconds timestamp precision
156    Nanoseconds,
157}
158
159impl TimestampPrecision {
160    fn api_short_name(&self) -> &str {
161        match self {
162            Self::Seconds => "s",
163            Self::Milliseconds => "ms",
164            Self::Microseconds => "us",
165            Self::Nanoseconds => "ns",
166        }
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use crate::models::DataPoint;
174    use futures::stream;
175    use mockito::mock;
176
177    #[tokio::test]
178    async fn writing_points() {
179        let org = "some-org";
180        let bucket = "some-bucket";
181        let token = "some-token";
182
183        let mock_server = mock(
184            "POST",
185            format!("/api/v2/write?bucket={}&org={}&precision=ns", bucket, org).as_str(),
186        )
187        .match_header("Authorization", format!("Token {}", token).as_str())
188        .match_body(
189            "\
190cpu,host=server01 usage=0.5
191cpu,host=server01,region=us-west usage=0.87
192",
193        )
194        .with_status(204)
195        .create();
196
197        let client = Client::new(mockito::server_url(), org, token);
198
199        let points = vec![
200            DataPoint::builder("cpu")
201                .tag("host", "server01")
202                .field("usage", 0.5)
203                .build()
204                .unwrap(),
205            DataPoint::builder("cpu")
206                .tag("host", "server01")
207                .tag("region", "us-west")
208                .field("usage", 0.87)
209                .build()
210                .unwrap(),
211        ];
212
213        // If the requests made are incorrect, Mockito returns status 501 and `write`
214        // will return an error, which causes the test to fail here instead of
215        // when we assert on mock_server. The error messages that Mockito
216        // provides are much clearer for explaining why a test failed than just
217        // that the server returned 501, so don't use `?` here.
218        let result = client.write(bucket, stream::iter(points)).await;
219        mock_server.assert();
220        assert!(result.is_ok());
221    }
222
223    #[tokio::test]
224    async fn writing_points_with_precision() {
225        let org = "some-org";
226        let bucket = "some-bucket";
227        let token = "some-token";
228
229        let mock_server = mock(
230            "POST",
231            format!("/api/v2/write?bucket={}&org={}&precision=s", bucket, org).as_str(),
232        )
233        .match_header("Authorization", format!("Token {}", token).as_str())
234        .match_body(
235            "\
236cpu,host=server01 usage=0.5 1671095854
237",
238        )
239        .with_status(204)
240        .create();
241
242        let client = Client::new(mockito::server_url(), org, token);
243
244        let point = DataPoint::builder("cpu")
245            .tag("host", "server01")
246            .field("usage", 0.5)
247            .timestamp(1671095854)
248            .build()
249            .unwrap();
250        let points = vec![point];
251
252        // If the requests made are incorrect, Mockito returns status 501 and `write`
253        // will return an error, which causes the test to fail here instead of
254        // when we assert on mock_server. The error messages that Mockito
255        // provides are much clearer for explaining why a test failed than just
256        // that the server returned 501, so don't use `?` here.
257        let result = client
258            .write_with_precision(bucket, stream::iter(points), TimestampPrecision::Seconds)
259            .await;
260        mock_server.assert();
261        assert!(result.is_ok());
262    }
263
264    #[tokio::test]
265    async fn status_code_correctly_interpreted() {
266        let org = "org";
267        let token = "token";
268        let bucket = "bucket";
269
270        let make_mock_server = |status| {
271            mock(
272                "POST",
273                format!("/api/v2/write?bucket={}&org={}&precision=ns", bucket, org).as_str(),
274            )
275            .with_status(status)
276            .create()
277        };
278
279        let write_with_status = |status| async move {
280            let mock_server = make_mock_server(status);
281            let client = Client::new(mockito::server_url(), org, token);
282            let points: Vec<DataPoint> = vec![];
283            let res = client.write(bucket, stream::iter(points)).await;
284            mock_server.assert();
285            res
286        };
287
288        // success status
289        assert!(write_with_status(204).await.is_ok());
290
291        // failing status
292        for status in [200, 201, 400, 401, 404, 413, 429, 500, 503] {
293            assert!(write_with_status(status).await.is_err());
294        }
295    }
296}
clia_influxdb2/api/write.rs

clia_influxdb2/api/
write.rs
