lora_executor/pull/
source.rs1use crate::errors::ExecResult;
2use crate::value::Row;
3
4pub trait RowSource {
11 fn next_row(&mut self) -> ExecResult<Option<Row>>;
13}
14
15pub fn drain<S: RowSource + ?Sized>(source: &mut S) -> ExecResult<Vec<Row>> {
17 let mut out = Vec::new();
18 while let Some(row) = source.next_row()? {
19 out.push(row);
20 }
21 Ok(out)
22}
23
24pub struct BufferedRowSource {
28 iter: std::vec::IntoIter<Row>,
29}
30
31impl BufferedRowSource {
32 pub fn new(rows: Vec<Row>) -> Self {
33 Self {
34 iter: rows.into_iter(),
35 }
36 }
37}
38
39impl RowSource for BufferedRowSource {
40 fn next_row(&mut self) -> ExecResult<Option<Row>> {
41 Ok(self.iter.next())
42 }
43}
44
45pub struct ArgumentSource {
48 yielded: bool,
49}
50
51impl ArgumentSource {
52 pub fn new() -> Self {
53 Self { yielded: false }
54 }
55}
56
57impl Default for ArgumentSource {
58 fn default() -> Self {
59 Self::new()
60 }
61}
62
63impl RowSource for ArgumentSource {
64 fn next_row(&mut self) -> ExecResult<Option<Row>> {
65 if self.yielded {
66 Ok(None)
67 } else {
68 self.yielded = true;
69 Ok(Some(Row::new()))
70 }
71 }
72}
73
74pub(crate) struct DeadlineSource<'a> {
80 inner: Box<dyn RowSource + 'a>,
81 deadline: web_time::Instant,
82 tick: u32,
83}
84
85impl<'a> DeadlineSource<'a> {
86 pub(crate) fn wrap(
87 inner: Box<dyn RowSource + 'a>,
88 deadline: Option<web_time::Instant>,
89 ) -> Box<dyn RowSource + 'a> {
90 match deadline {
91 Some(deadline) => Box::new(Self {
92 inner,
93 deadline,
94 tick: 0,
95 }),
96 None => inner,
97 }
98 }
99}
100
101impl RowSource for DeadlineSource<'_> {
102 fn next_row(&mut self) -> ExecResult<Option<Row>> {
103 self.tick = self.tick.wrapping_add(1);
104 if self.tick.is_multiple_of(64) && crate::cancel::deadline_reached(self.deadline) {
105 return Err(crate::errors::ExecutorError::QueryTimeout);
106 }
107 self.inner.next_row()
108 }
109}