Skip to main content

hematite/sql/
script.rs

1//! Token-aware SQL script splitting and stepped execution.
2
3use std::collections::VecDeque;
4
5use crate::error::Result;
6use crate::parser::lexer::Token;
7use crate::parser::{Lexer, Parser};
8
9use super::connection::Connection;
10use super::result::ExecutedStatement;
11
12pub(crate) fn split_script_tokens(sql: &str) -> Result<Vec<Vec<Token>>> {
13    Ok(split_script_state(sql, true)?.0)
14}
15
16pub fn script_is_complete(sql: &str) -> Result<bool> {
17    let (statements, has_incomplete_tail) = split_script_state(sql, false)?;
18    Ok(!statements.is_empty() && !has_incomplete_tail)
19}
20
21fn split_script_state(
22    sql: &str,
23    append_trailing_statement: bool,
24) -> Result<(Vec<Vec<Token>>, bool)> {
25    let mut lexer = Lexer::new(sql.to_string());
26    lexer.tokenize()?;
27
28    let mut statements = Vec::new();
29    let mut current_tokens = Vec::new();
30
31    for token in lexer.get_tokens().iter().cloned() {
32        let is_semicolon = matches!(token, Token::Semicolon);
33        current_tokens.push(token);
34
35        if is_semicolon {
36            if contains_statement_tokens(&current_tokens) {
37                statements.push(current_tokens);
38            }
39            current_tokens = Vec::new();
40        }
41    }
42
43    if contains_statement_tokens(&current_tokens) {
44        if append_trailing_statement {
45            current_tokens.push(Token::Semicolon);
46            statements.push(current_tokens);
47            return Ok((statements, false));
48        }
49        return Ok((statements, true));
50    }
51
52    Ok((statements, false))
53}
54
55fn contains_statement_tokens(tokens: &[Token]) -> bool {
56    tokens
57        .iter()
58        .any(|token| !matches!(token, Token::Semicolon))
59}
60
61pub struct ScriptIter<'a> {
62    connection: &'a mut Connection,
63    statements: VecDeque<Vec<Token>>,
64}
65
66impl<'a> ScriptIter<'a> {
67    pub(crate) fn new(connection: &'a mut Connection, statements: Vec<Vec<Token>>) -> Self {
68        Self {
69            connection,
70            statements: statements.into(),
71        }
72    }
73}
74
75impl Iterator for ScriptIter<'_> {
76    type Item = Result<ExecutedStatement>;
77
78    fn next(&mut self) -> Option<Self::Item> {
79        let tokens = self.statements.pop_front()?;
80        let mut parser = Parser::new(tokens);
81        Some(
82            parser
83                .parse()
84                .and_then(|statement| self.connection.execute_statement_result(statement)),
85        )
86    }
87}