sofie_db/
stream.rs

1use std::collections::VecDeque;
2
3use sofie_core::{Cast, Error, Result, SqlType, SqlValue, Uuid};
4
5use crate::condition::Kind;
6use crate::statement::{Criteria, SortKey};
7use crate::{Read, Selected, Source};
8
9const BATCH: usize = 500;
10
11/// One short query per batch, so no cursor or transaction stays open between them.
12pub struct Stream<T: Source> {
13    criteria: Criteria,
14    read: Read,
15    batch: VecDeque<T>,
16    after: Option<Uuid>,
17    exhausted: bool,
18}
19
20impl<T: Source> Stream<T> {
21    pub(crate) fn new(criteria: Criteria, read: &Read) -> Result<Self> {
22        if !criteria.sorts.is_empty() || criteria.limit.is_some() {
23            return Err(Error::internal_message(
24                "a stream reads every matching record in id order; remove its `order_by`, \
25                 `latest` and `limit`",
26            ));
27        }
28        Ok(Self {
29            criteria,
30            read: read.clone(),
31            batch: VecDeque::new(),
32            after: None,
33            exhausted: false,
34        })
35    }
36
37    pub async fn next(&mut self) -> Result<Option<T>> {
38        if self.batch.is_empty() && !self.exhausted {
39            self.read_batch().await?;
40        }
41        Ok(self.batch.pop_front())
42    }
43
44    async fn read_batch(&mut self) -> Result<()> {
45        if self.after.is_none() {
46            self.criteria = std::mem::take(&mut self.criteria)
47                .matched(&self.read)
48                .await?;
49        }
50        let mut criteria = self.criteria.clone();
51        if let Some(after) = self.after {
52            criteria.conditions.push(Kind::Compare {
53                columns: vec!["id".to_owned()],
54                operator: ">",
55                values: vec![SqlValue::Uuid(after)],
56            });
57        }
58        criteria.sorts.push(SortKey::id(false));
59        criteria.limit = u64::try_from(BATCH).ok();
60        let mut columns = vec![Selected::column("id")];
61        columns.extend(T::columns());
62        let (sql, params) =
63            criteria.select_sql::<T::Base>(self.read.dialect(), self.read.owner(), &columns);
64        let mut types = vec![Some(SqlType::Uuid)];
65        types.extend(T::column_types());
66        let rows = self
67            .read
68            .decoded(&sql, &params, BATCH, types, |row| {
69                Ok((Uuid::decode(row)?, T::decode(row)?))
70            })
71            .await?;
72        self.exhausted = rows.len() < BATCH;
73        for (id, item) in rows {
74            self.after = Some(id);
75            self.batch.push_back(item);
76        }
77        Ok(())
78    }
79}
sofie_db/stream.rs

sofie_db/
stream.rs
