systemprompt_database/services/
executor.rs

1//! SQL batch and statement-by-statement execution helpers.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use super::database::Database;
7use super::provider::DatabaseProvider;
8use crate::error::{DatabaseResult, RepositoryError};
9use crate::models::QueryResult;
10
11#[derive(Debug, Copy, Clone)]
12pub struct SqlExecutor;
13
14enum SplitState {
15    Normal,
16    SingleQuote,
17    DollarQuote(String),
18    LineComment,
19    BlockComment(u32),
20}
21
22fn dollar_tag_end(bytes: &[u8], start: usize) -> Option<usize> {
23    debug_assert_eq!(bytes[start], b'$');
24    let mut j = start + 1;
25    while j < bytes.len() {
26        let c = bytes[j];
27        if c == b'$' {
28            return Some(j);
29        }
30        if !(c.is_ascii_alphanumeric() || c == b'_') {
31            return None;
32        }
33        j += 1;
34    }
35    None
36}
37
38struct Splitter<'a> {
39    sql: &'a str,
40    bytes: &'a [u8],
41    i: usize,
42    start: usize,
43    has_content: bool,
44    statements: Vec<String>,
45}
46
47impl<'a> Splitter<'a> {
48    const fn new(sql: &'a str) -> Self {
49        Self {
50            sql,
51            bytes: sql.as_bytes(),
52            i: 0,
53            start: 0,
54            has_content: false,
55            statements: Vec::new(),
56        }
57    }
58
59    fn emit(&mut self, end: usize) {
60        if self.has_content {
61            let stmt = self.sql[self.start..end].trim();
62            if !stmt.is_empty() {
63                self.statements.push(stmt.to_owned());
64            }
65        }
66        self.has_content = false;
67    }
68
69    fn step_normal(&mut self) -> SplitState {
70        match self.bytes[self.i] {
71            b'\'' => {
72                self.has_content = true;
73                self.i += 1;
74                SplitState::SingleQuote
75            },
76            b'-' if self.bytes.get(self.i + 1) == Some(&b'-') => {
77                self.i += 2;
78                SplitState::LineComment
79            },
80            b'/' if self.bytes.get(self.i + 1) == Some(&b'*') => {
81                self.i += 2;
82                SplitState::BlockComment(1)
83            },
84            b'$' => {
85                self.has_content = true;
86                if let Some(tag_end) = dollar_tag_end(self.bytes, self.i) {
87                    let tag = self.sql[self.i..=tag_end].to_string();
88                    self.i = tag_end + 1;
89                    SplitState::DollarQuote(tag)
90                } else {
91                    self.i += 1;
92                    SplitState::Normal
93                }
94            },
95            b';' => {
96                self.emit(self.i);
97                self.i += 1;
98                self.start = self.i;
99                SplitState::Normal
100            },
101            b => {
102                if !b.is_ascii_whitespace() {
103                    self.has_content = true;
104                }
105                self.i += 1;
106                SplitState::Normal
107            },
108        }
109    }
110
111    fn step_single_quote(&mut self) -> SplitState {
112        if self.bytes[self.i] == b'\'' {
113            if self.bytes.get(self.i + 1) == Some(&b'\'') {
114                self.i += 2;
115                SplitState::SingleQuote
116            } else {
117                self.i += 1;
118                SplitState::Normal
119            }
120        } else {
121            self.i += 1;
122            SplitState::SingleQuote
123        }
124    }
125
126    fn step_dollar_quote(&mut self, tag: String) -> SplitState {
127        let tag_bytes = tag.as_bytes();
128        if self.i + tag_bytes.len() <= self.bytes.len()
129            && self.bytes[self.i..self.i + tag_bytes.len()] == *tag_bytes
130        {
131            self.i += tag_bytes.len();
132            SplitState::Normal
133        } else {
134            self.i += 1;
135            SplitState::DollarQuote(tag)
136        }
137    }
138
139    const fn step_line_comment(&mut self) -> SplitState {
140        let next = if self.bytes[self.i] == b'\n' {
141            SplitState::Normal
142        } else {
143            SplitState::LineComment
144        };
145        self.i += 1;
146        next
147    }
148
149    fn step_block_comment(&mut self, depth: u32) -> SplitState {
150        if self.bytes[self.i] == b'/' && self.bytes.get(self.i + 1) == Some(&b'*') {
151            self.i += 2;
152            SplitState::BlockComment(depth + 1)
153        } else if self.bytes[self.i] == b'*' && self.bytes.get(self.i + 1) == Some(&b'/') {
154            self.i += 2;
155            if depth == 1 {
156                SplitState::Normal
157            } else {
158                SplitState::BlockComment(depth - 1)
159            }
160        } else {
161            self.i += 1;
162            SplitState::BlockComment(depth)
163        }
164    }
165
166    fn run(mut self) -> DatabaseResult<Vec<String>> {
167        let mut state = SplitState::Normal;
168        while self.i < self.bytes.len() {
169            state = match state {
170                SplitState::Normal => self.step_normal(),
171                SplitState::SingleQuote => self.step_single_quote(),
172                SplitState::DollarQuote(tag) => self.step_dollar_quote(tag),
173                SplitState::LineComment => self.step_line_comment(),
174                SplitState::BlockComment(depth) => self.step_block_comment(depth),
175            };
176        }
177
178        match state {
179            SplitState::Normal | SplitState::LineComment => {
180                let end = self.sql.len();
181                self.emit(end);
182                Ok(self.statements)
183            },
184            SplitState::SingleQuote => Err(RepositoryError::Internal(
185                "Unterminated string literal in SQL".into(),
186            )),
187            SplitState::DollarQuote(tag) => Err(RepositoryError::Internal(format!(
188                "Unterminated dollar-quoted string: {tag}"
189            ))),
190            SplitState::BlockComment(_) => Err(RepositoryError::Internal(
191                "Unterminated block comment in SQL".into(),
192            )),
193        }
194    }
195}
196
197impl SqlExecutor {
198    pub async fn execute_statements(db: &Database, sql: &str) -> DatabaseResult<()> {
199        db.execute_batch(sql).await.map_err(|e| {
200            RepositoryError::Internal(format!("Failed to execute SQL statements: {e}"))
201        })
202    }
203
204    pub async fn execute_statements_parsed(
205        db: &dyn DatabaseProvider,
206        sql: &str,
207    ) -> DatabaseResult<()> {
208        let statements = Self::parse_sql_statements(sql)?;
209
210        for statement in statements {
211            db.execute_raw(&statement).await.map_err(|e| {
212                RepositoryError::Internal(format!(
213                    "Failed to execute SQL statement: {statement}: {e}"
214                ))
215            })?;
216        }
217
218        Ok(())
219    }
220
221    pub fn parse_sql_statements(sql: &str) -> DatabaseResult<Vec<String>> {
222        Splitter::new(sql).run()
223    }
224
225    pub async fn execute_query(db: &Database, query: &str) -> DatabaseResult<QueryResult> {
226        db.query_raw(&query)
227            .await
228            .map_err(|e| RepositoryError::QueryExecution(Box::new(e)))
229    }
230
231    pub async fn execute_file(db: &Database, file_path: &str) -> DatabaseResult<()> {
232        let sql = std::fs::read_to_string(file_path).map_err(|e| {
233            RepositoryError::Internal(format!("Failed to read SQL file: {file_path}: {e}"))
234        })?;
235        Self::execute_statements(db, &sql).await
236    }
237
238    pub async fn table_exists(db: &Database, table_name: &str) -> DatabaseResult<bool> {
239        let result = db
240            .query_raw_with(
241                &"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = \
242                  'public' AND table_name = $1) as exists",
243                &[&table_name],
244            )
245            .await?;
246
247        result
248            .first()
249            .and_then(|row| row.get("exists"))
250            .and_then(serde_json::Value::as_bool)
251            .ok_or_else(|| RepositoryError::Internal("Failed to check table existence".to_owned()))
252    }
253
254    pub async fn column_exists(
255        db: &Database,
256        table_name: &str,
257        column_name: &str,
258    ) -> DatabaseResult<bool> {
259        let result = db
260            .query_raw_with(
261                &"SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema = \
262                  'public' AND table_name = $1 AND column_name = $2) as exists",
263                &[&table_name, &column_name],
264            )
265            .await?;
266
267        result
268            .first()
269            .and_then(|row| row.get("exists"))
270            .and_then(serde_json::Value::as_bool)
271            .ok_or_else(|| RepositoryError::Internal("Failed to check column existence".to_owned()))
272    }
273}
systemprompt_database/services/executor.rs

systemprompt_database/services/
executor.rs
