revolt_database/util/
chunked.rs

1#[cfg(feature = "mongodb")]
2use ::mongodb::{ClientSession, SessionCursor};
3use revolt_result::{Result, ToRevoltError};
4use serde::Deserialize;
5
6#[derive(Debug)]
7#[allow(clippy::large_enum_variant)]
8pub enum ChunkedDatabaseGenerator<T> {
9    #[cfg(feature = "mongodb")]
10    MongoDb {
11        session: ClientSession,
12        cursor: SessionCursor<T>,
13    },
14
15    Reference {
16        offset: usize,
17        data: Vec<T>,
18    },
19}
20
21impl<T: for<'d> Deserialize<'d> + Clone> ChunkedDatabaseGenerator<T> {
22    #[cfg(feature = "mongodb")]
23    pub fn new_mongo(session: ClientSession, cursor: SessionCursor<T>) -> Self {
24        Self::MongoDb { session, cursor }
25    }
26
27    pub fn new_reference(data: Vec<T>) -> Self {
28        Self::Reference { offset: 0, data }
29    }
30
31    pub async fn next(&mut self) -> Result<Option<T>> {
32        match self {
33            #[cfg(feature = "mongodb")]
34            Self::MongoDb { session, cursor } => {
35                cursor.next(session).await.transpose().to_internal_error()
36            }
37            Self::Reference { offset, data } => {
38                if let Some(value) = data.get(*offset) {
39                    *offset += 1;
40                    Ok(Some(value.clone()))
41                } else {
42                    Ok(None)
43                }
44            }
45        }
46    }
47
48    pub async fn next_n(&mut self, n: usize) -> Result<Option<Vec<T>>> {
49        let mut docs = Vec::new();
50
51        while docs.len() < n {
52            if let Some(doc) = self.next().await? {
53                docs.push(doc);
54            } else if docs.is_empty() {
55                return Ok(None);
56            } else {
57                break;
58            }
59        }
60
61        Ok(Some(docs))
62    }
63}
revolt_database/util/chunked.rs

revolt_database/util/
chunked.rs
