reduct_rs/record/
query.rs

1// Copyright 2023-2026 ReductStore
2// This Source Code Form is subject to the terms of the Mozilla Public
3//    License, v. 2.0. If a copy of the MPL was not distributed with this
4//    file, You can obtain one at https://mozilla.org/MPL/2.0/.
5
6use crate::http_client::{map_error, HttpClient};
7use crate::record::{from_system_time, Record};
8use crate::RecordStream;
9use async_channel::{unbounded, Receiver};
10use async_stream::stream;
11use bytes::Bytes;
12use bytes::BytesMut;
13use futures::Stream;
14use futures_util::{pin_mut, StreamExt};
15use reduct_base::batch::v2::{parse_batched_headers, EntryRecordHeader};
16use reduct_base::batch::{parse_batched_header, sort_headers_by_time, RecordHeader};
17use reduct_base::error::ErrorCode::Unknown;
18use reduct_base::error::{ErrorCode, ReductError};
19use reduct_base::msg::entry_api::{QueryEntry, QueryInfo, QueryType, RemoveQueryInfo};
20use reqwest::header::{HeaderMap, HeaderValue};
21use reqwest::Method;
22use serde_json::Value;
23use std::pin::Pin;
24use std::sync::Arc;
25use std::time::{Duration, SystemTime};
26
27type QueryStream = Pin<Box<dyn Stream<Item = Result<Record, ReductError>> + Send>>;
28
29/// Builder for a query request.
30pub struct QueryBuilder {
31    query: QueryEntry,
32
33    bucket: String,
34    entries: Vec<String>,
35    client: Arc<HttpClient>,
36}
37
38impl QueryBuilder {
39    pub(crate) fn new(bucket: String, entries: Vec<String>, client: Arc<HttpClient>) -> Self {
40        Self {
41            query: QueryEntry::default(),
42            bucket,
43            entries,
44            client,
45        }
46    }
47
48    /// Set the start time of the query.
49    pub fn start(mut self, time: SystemTime) -> Self {
50        self.query.start = Some(from_system_time(time));
51        self
52    }
53
54    /// Set the start time of the query as a unix timestamp in microseconds.
55    pub fn start_us(mut self, time_us: u64) -> Self {
56        self.query.start = Some(time_us);
57        self
58    }
59
60    /// Set the end time of the query.
61    pub fn stop(mut self, time: SystemTime) -> Self {
62        self.query.stop = Some(from_system_time(time));
63        self
64    }
65
66    /// Set the end time of the query as a unix timestamp in microseconds.
67    pub fn stop_us(mut self, time_us: u64) -> Self {
68        self.query.stop = Some(time_us);
69        self
70    }
71
72    /// Set the condition for the query.
73    pub fn when(mut self, condition: Value) -> Self {
74        self.query.when = Some(condition);
75        self
76    }
77
78    /// Set the query to be strict.
79    /// If the query is strict, the query will return an error if any of the conditions are invalid.
80    /// default: false
81    pub fn strict(mut self, strict: bool) -> Self {
82        self.query.strict = Some(strict);
83        self
84    }
85
86    /// Set extension parameters for the query.
87    /// This is a JSON object that will be passed to extensions on the server side.
88    pub fn ext(mut self, ext: Value) -> Self {
89        self.query.ext = Some(ext);
90        self
91    }
92
93    /// Set TTL for the query.
94    pub fn ttl(mut self, ttl: Duration) -> Self {
95        self.query.ttl = Some(ttl.as_secs());
96        self
97    }
98
99    /// Set the query to be continuous.
100    pub fn continuous(mut self) -> Self {
101        self.query.continuous = Some(true);
102        self
103    }
104
105    /// Set the query to head only.
106    /// default: false
107    pub fn head_only(mut self, head_only: bool) -> Self {
108        self.query.only_metadata = Some(head_only);
109        self
110    }
111
112    /// Send the query request.
113    pub async fn send(
114        self,
115    ) -> Result<impl Stream<Item = Result<Record, ReductError>>, ReductError> {
116        if self.entries.len() == 1 {
117            self.query_v1().await
118        } else {
119            if let Some(version) = self.client.get_api_version().await {
120                if version.1 < 18 {
121                    return Err(ReductError::new(
122                        ErrorCode::InvalidRequest,
123                        "Multi-entry queries are not supported in API versions below v1.18",
124                    ));
125                }
126            }
127
128            self.query_v2().await
129        }
130    }
131
132    async fn query_v1(mut self) -> Result<QueryStream, ReductError> {
133        self.query.query_type = QueryType::Query;
134        let entry = self.entries.first().cloned().unwrap();
135
136        let response = self
137            .client
138            .send_and_receive_json::<QueryEntry, QueryInfo>(
139                Method::POST,
140                &format!("/b/{}/{}/q", self.bucket, entry),
141                Some(self.query.clone()),
142            )
143            .await?;
144
145        let head_only = self.query.only_metadata.as_ref().unwrap_or(&false).clone();
146
147        Ok(Box::pin(stream! {
148            let mut last = false;
149            while !last {
150                let method = if head_only { Method::HEAD } else { Method::GET };
151                let request = self.client.request(
152                    method,
153                    &format!("/b/{}/{}/batch?q={}", self.bucket, entry, response.id),
154                );
155                let response = self.client.send_request(request).await?;
156
157                if response.status() == reqwest::StatusCode::NO_CONTENT {
158                    break;
159                }
160
161                let headers = response.headers().clone();
162
163
164                let (tx, rx) = unbounded();
165                tokio::spawn(async move {
166                    let mut stream = response.bytes_stream();
167                    while let Some(bytes) = stream.next().await {
168                        if let Err(_) = tx.send(bytes).await {
169                            break;
170                        }
171                    }
172                });
173
174                let stream = parse_batched_records(&entry, headers, rx, head_only).await?;
175                pin_mut!(stream);
176                while let Some(record) = stream.next().await {
177                    let record = record?;
178                    last = record.1;
179                    yield Ok(record.0);
180                }
181            }
182        }))
183    }
184
185    async fn query_v2(mut self) -> Result<QueryStream, ReductError> {
186        self.query.query_type = QueryType::Query;
187        self.query.entries = Some(self.entries.clone());
188        let response = self
189            .client
190            .send_and_receive_json::<QueryEntry, QueryInfo>(
191                Method::POST,
192                &format!("/io/{}/q", self.bucket),
193                Some(self.query.clone()),
194            )
195            .await?;
196
197        let head_only = self.query.only_metadata.as_ref().unwrap_or(&false).clone();
198
199        Ok(Box::pin(stream! {
200            let mut last = false;
201            while !last {
202                let method = if head_only { Method::HEAD } else { Method::GET };
203                let request = self
204                    .client
205                    .request(method, &format!("/io/{}/read", self.bucket))
206                    .header("x-reduct-query-id", response.id.to_string());
207                let response = self.client.send_request(request).await?;
208
209                if response.status() == reqwest::StatusCode::NO_CONTENT {
210                    break;
211                }
212
213                let headers = response.headers().clone();
214
215                let (tx, rx) = unbounded();
216                tokio::spawn(async move {
217                    let mut stream = response.bytes_stream();
218                    while let Some(bytes) = stream.next().await {
219                        if let Err(_) = tx.send(bytes).await {
220                            break;
221                        }
222                    }
223                });
224
225                let stream = parse_batched_records_v2(headers, rx, head_only).await?;
226                pin_mut!(stream);
227                while let Some(record) = stream.next().await {
228                    let record = record?;
229                    last = record.1;
230                    yield Ok(record.0);
231                }
232            }
233        }))
234    }
235}
236
237/**
238 * Builder for a remove query request.
239 */
240pub struct RemoveQueryBuilder {
241    query: QueryEntry,
242
243    bucket: String,
244    entries: Vec<String>,
245    client: Arc<HttpClient>,
246}
247
248impl RemoveQueryBuilder {
249    pub(crate) fn new(bucket: String, entries: Vec<String>, client: Arc<HttpClient>) -> Self {
250        Self {
251            query: QueryEntry::default(),
252            bucket,
253            entries,
254            client,
255        }
256    }
257
258    /// Set the start time of the query.
259    pub fn start(mut self, time: SystemTime) -> Self {
260        self.query.start = Some(from_system_time(time));
261        self
262    }
263
264    /// Set the start time of the query as a unix timestamp in microseconds.
265    pub fn start_us(mut self, time_us: u64) -> Self {
266        self.query.start = Some(time_us);
267        self
268    }
269
270    /// Set the end time of the query.
271    pub fn stop(mut self, time: SystemTime) -> Self {
272        self.query.stop = Some(from_system_time(time));
273        self
274    }
275
276    /// Set the end time of the query as a unix timestamp in microseconds.
277    pub fn stop_us(mut self, time_us: u64) -> Self {
278        self.query.stop = Some(time_us);
279        self
280    }
281
282    /// Set the condition for the query.
283    /// This will remove all records that match the condition.
284    /// This is a destructive operation.
285    pub fn when(mut self, condition: Value) -> Self {
286        self.query.when = Some(condition);
287        self
288    }
289
290    /// Set the query to be strict.
291    /// If the query is strict, the query will return an error if any of the conditions are invalid.
292    /// default: false
293    pub fn strict(mut self, strict: bool) -> Self {
294        self.query.strict = Some(strict);
295        self
296    }
297
298    /// Send the remove query request.
299    /// This will remove all records that match the query.
300    /// This is a destructive operation.
301    ///
302    /// # Returns
303    ///
304    /// * `Result<u64, ReductError>` - The number of records removed.
305    pub async fn send(mut self) -> Result<u64, ReductError> {
306        self.query.query_type = QueryType::Remove;
307        if self.entries.len() == 1 {
308            let entry = self.entries.first().cloned().unwrap();
309            let response = self
310                .client
311                .send_and_receive_json::<QueryEntry, RemoveQueryInfo>(
312                    Method::POST,
313                    &format!("/b/{}/{}/q", self.bucket, entry),
314                    Some(self.query.clone()),
315                )
316                .await?;
317
318            Ok(response.removed_records)
319        } else {
320            if let Some(version) = self.client.get_api_version().await {
321                if version.1 < 18 {
322                    return Err(ReductError::new(
323                        ErrorCode::InvalidRequest,
324                        "Multi-entry remove queries are not supported in API versions below v1.18",
325                    ));
326                }
327            }
328
329            self.query.entries = Some(self.entries.clone());
330            let response = self
331                .client
332                .send_and_receive_json::<QueryEntry, RemoveQueryInfo>(
333                    Method::POST,
334                    &format!("/io/{}/q", self.bucket),
335                    Some(self.query.clone()),
336                )
337                .await?;
338
339            Ok(response.removed_records)
340        }
341    }
342}
343
344async fn parse_batched_records(
345    queried_entry: &str,
346    headers: HeaderMap,
347    rx: Receiver<Result<Bytes, reqwest::Error>>,
348    head_only: bool,
349) -> Result<impl Stream<Item = Result<(Record, bool), ReductError>>, ReductError> {
350    let sorted_records = sort_headers_by_time(&headers)?;
351    let last = headers.get("x-reduct-last") == Some(&HeaderValue::from_str("true").unwrap());
352    let mut records = Vec::with_capacity(sorted_records.len());
353
354    for (timestamp, value) in sorted_records {
355        let RecordHeader {
356            content_length,
357            content_type,
358            labels,
359        } = parse_batched_header(value.to_str().unwrap()).unwrap();
360        records.push((
361            timestamp,
362            queried_entry.to_string(),
363            RecordHeader {
364                content_length,
365                content_type,
366                labels,
367            },
368        ));
369    }
370
371    parse_batched_records_from_headers(records, last, rx, head_only).await
372}
373
374async fn parse_batched_records_v2(
375    headers: HeaderMap,
376    rx: Receiver<Result<Bytes, reqwest::Error>>,
377    head_only: bool,
378) -> Result<impl Stream<Item = Result<(Record, bool), ReductError>>, ReductError> {
379    let sorted_records = parse_batched_headers(&headers)?;
380    let last = headers.get("x-reduct-last") == Some(&HeaderValue::from_str("true").unwrap());
381    let records = sorted_records
382        .into_iter()
383        .map(
384            |EntryRecordHeader {
385                 timestamp,
386                 entry,
387                 header,
388                 ..
389             }| (timestamp, entry, header),
390        )
391        .collect::<Vec<_>>();
392
393    parse_batched_records_from_headers(records, last, rx, head_only).await
394}
395
396async fn parse_batched_records_from_headers(
397    records: Vec<(u64, String, RecordHeader)>,
398    last: bool,
399    rx: Receiver<Result<Bytes, reqwest::Error>>,
400    head_only: bool,
401) -> Result<impl Stream<Item = Result<(Record, bool), ReductError>>, ReductError> {
402    let records_total = records.len();
403    let mut records_count = 0;
404
405    let unwrap_byte = |bytes: Result<Bytes, reqwest::Error>| match bytes {
406        Ok(b) => Ok(b),
407        Err(err) => {
408            if let Some(status) = err.status() {
409                Err(ReductError::new(
410                    ErrorCode::try_from(status.as_u16() as i16).unwrap_or(Unknown),
411                    &err.to_string(),
412                ))
413            } else {
414                Err(map_error(err))
415            }
416        }
417    };
418
419    Ok(stream! {
420        let mut rest_data = BytesMut::new();
421
422        for (timestamp, entry, header) in records.into_iter() {
423            let RecordHeader { content_length, content_type, labels } = header;
424            records_count += 1;
425
426            let data: Option<RecordStream> = if head_only {
427                None
428            } else if records_count == records_total {
429                let first_chunk: Bytes = rest_data.clone().into();
430                let rx = rx.clone();
431
432                Some(Box::pin(stream! {
433                    yield Ok(first_chunk);
434                    while let Ok(bytes) = rx.recv().await {
435                        yield unwrap_byte(bytes);
436                    }
437                }))
438            } else {
439                let mut data = rest_data.clone();
440                while let Ok(bytes) = rx.recv().await {
441                    data.extend_from_slice(&unwrap_byte(bytes)?);
442                    if data.len() >= content_length  as usize {
443                        break;
444                    }
445                }
446
447                rest_data = data.split_off(content_length as usize);
448                data.truncate(content_length as usize);
449
450                Some(Box::pin(stream! {
451                    yield Ok(data.into());
452                }))
453            };
454
455            yield Ok((Record {
456                timestamp,
457                entry,
458                labels,
459                content_type,
460                content_length,
461                data
462            }, last));
463        }
464    })
465}
reduct_rs/record/query.rs

reduct_rs/record/
query.rs
