restq/
multi_stmt.rs

1use crate::StmtData;
2use std::io::{BufRead, BufReader, Cursor, Read, Seek, SeekFrom};
3
4/// A statement iterator, lazily parse statement with data
5pub struct MultiStatement<R> {
6    content: BufReader<R>,
7}
8
9impl<R> MultiStatement<R>
10where
11    R: Read + Send + Sync + Seek,
12{
13    pub fn from_reader(reader: R) -> Self {
14        MultiStatement {
15            content: BufReader::new(reader),
16        }
17    }
18
19    pub fn statement_iter(self) -> StatementIter<R> {
20        StatementIter {
21            content: self.content,
22        }
23    }
24}
25
26pub struct StatementIter<R> {
27    content: BufReader<R>,
28}
29
30impl<R> StatementIter<R>
31where
32    R: Read + Send + Sync + Seek,
33{
34    pub fn new(content: BufReader<R>) -> Self {
35        StatementIter { content }
36    }
37
38    fn ignore_whitespace(&mut self) {
39        let mut buffer = String::new();
40        while let Ok(n) = self.content.read_line(&mut buffer) {
41            if buffer.trim().is_empty() {
42                //ignoring blank lines
43            } else {
44                // seek to the last non blank character than break the loop
45                self.content
46                    .seek(SeekFrom::Current(-(n as i64)))
47                    .expect("must seek");
48                break;
49            }
50            if n == 0 {
51                break;
52            }
53        }
54    }
55}
56
57impl<R> Iterator for StatementIter<R>
58where
59    R: Read + Send + Sync + Seek,
60{
61    type Item = StmtData<Cursor<Vec<u8>>>;
62
63    fn next(&mut self) -> Option<Self::Item> {
64        self.ignore_whitespace();
65        let mut buffer = vec![];
66        while let Ok(n) = self.content.read_until(b'\n', &mut buffer) {
67            let last_char = buffer.iter().last();
68            if (n == 1 && last_char == Some(&b'\n')) || n == 0 {
69                if !buffer.is_empty() {
70                    let stmt = StmtData::from_reader(Cursor::new(buffer))
71                        .expect("must not error");
72                    return Some(stmt);
73                } else {
74                    return None;
75                }
76            }
77            if n == 0 {
78                break;
79            }
80        }
81        None
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::ast::*;
89
90    #[test]
91    fn test_multi_statement() {
92        let data = r#"PUT /+category{*category_id:s32,name:text,description:text?,slug:text?,topic:i32?,created_at:utc,updated:utc}
931,Staff,staff,Private categories for staff discussion. Topics are only visible to admin and moderators
942,Technology,tecnology,Anything related to technology
95
96PUT /+topic{*topic:s32,title:text,excerpt:text?,created_at:utc(now()),updated_at:utc?}
971,About Euphorum,1
983,Topic3,3
992,Welcome to Euphorum,2"#;
100
101        let ms = MultiStatement::from_reader(Cursor::new(data.as_bytes()));
102        let mut iter = ms.statement_iter();
103        let stmt1 = iter.next().expect("must have a next");
104        if let Statement::Create(create_cat) = stmt1.statement() {
105            println!("create1: {}", create_cat);
106        }
107        let data1 = stmt1.rows_iter().expect("must have csv rows");
108        let all_data1 = data1.collect::<Vec<_>>();
109        assert_eq!(all_data1.len(), 2);
110
111        let stmt2 = iter.next().expect("must have a next");
112        if let Statement::Create(create_topic) = stmt2.statement() {
113            println!("create2: {}", create_topic);
114        }
115        let data2 = stmt2.rows_iter().expect("must have csv rows");
116        let all_data2 = dbg!(data2.collect::<Vec<_>>());
117        assert_eq!(all_data2.len(), 3);
118    }
119
120    #[test]
121    fn test_multi_statement_with_multiple_empty_lines() {
122        let data = r#"
123
124PUT /+category{*category_id:s32,name:text,description:text?,slug:text?,topic:i32?,created_at:utc,updated:utc}
1251,Staff,staff,Private categories for staff discussion. Topics are only visible to admin and moderators
1262,Technology,tecnology,Anything related to technology
127
128
129
130PUT /+topic{*topic:s32,title:text,excerpt:text?,created_at:utc(now()),updated_at:utc?}
1311,About Euphorum,1
1323,Topic3,3
1332,Welcome to Euphorum,2"#;
134
135        let ms = MultiStatement::from_reader(Cursor::new(data.as_bytes()));
136        let mut iter = ms.statement_iter();
137        let stmt1 = iter.next().expect("must have a next");
138        if let Statement::Create(create_cat) = stmt1.statement() {
139            println!("create1: {}", create_cat);
140        }
141        let data1 = stmt1.rows_iter().expect("must have csv rows");
142        let all_data1 = data1.collect::<Vec<_>>();
143        assert_eq!(all_data1.len(), 2);
144
145        let stmt2 = iter.next().expect("must have a next");
146        if let Statement::Create(create_topic) = stmt2.statement() {
147            println!("create2: {}", create_topic);
148        }
149        let data2 = stmt2.rows_iter().expect("must have csv rows");
150        let all_data2 = dbg!(data2.collect::<Vec<_>>());
151        assert_eq!(all_data2.len(), 3);
152    }
153}