Skip to main content

arcature_data/lint/
statement.rs

1//! SQL statement tokenization: split a script into [`SqlStatement`]s,
2//! respecting string literals, quoted identifiers, and comments.
3//!
4//! The splitter is a **robust, conservative scanner** — not a full SQL
5//! parser. It walks the input byte by byte, skipping line comments (`--`)
6//! and block comments (`/* */`), tracking single-quoted string literals
7//! (with `''` escapes) and double-quoted identifiers, and splitting on
8//! top-level `;`. It never panics on hostile or malformed input: an
9//! unterminated string or comment simply ends at end-of-input, and the
10//! accumulated text becomes the final statement.
11
12/// A single SQL statement, normalized for classification.
13///
14/// Construct via [`SqlStatement::split`] (a script into statements) or
15/// [`SqlStatement::single`] (one statement). The `sql` field carries the
16/// normalized text the classifier scans; comments are stripped, but quoted
17/// literals and identifiers are preserved so keyword detection does not fire
18/// inside a string.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct SqlStatement {
21    /// The normalized statement text (comments stripped; statements trimmed).
22    pub sql: String,
23}
24
25impl SqlStatement {
26    /// Wrap a single already-split statement.
27    #[must_use]
28    pub fn single(sql: impl Into<String>) -> Self {
29        let sql = sql.into();
30        Self {
31            sql: sql.trim().to_owned(),
32        }
33    }
34
35    /// Split a SQL script into statements. Whitespace-only statements are
36    /// dropped. Never panics.
37    #[must_use]
38    pub fn split(script: &str) -> Vec<Self> {
39        let mut statements = Vec::new();
40        let mut current = String::new();
41        let bytes = script.as_bytes();
42        let mut index = 0;
43        while index < bytes.len() {
44            let rest = &script[index..];
45            // Line comment: `--` to end of line.
46            if rest.starts_with("--") {
47                let end = rest.find('\n').map(|n| n + 1).unwrap_or(rest.len());
48                index += end;
49                continue;
50            }
51            // Block comment: `/* ... */` (nesting not supported by PostgreSQL
52            // block comments, so a simple non-nesting scan is correct).
53            if rest.starts_with("/*") {
54                let end = rest.find("*/").map(|n| n + 2).unwrap_or(rest.len());
55                index += end;
56                continue;
57            }
58            // Single-quoted string literal: `'...'` with `''` escapes.
59            let bytes_rest = &bytes[index..];
60            if bytes_rest.first().is_some_and(|&b| b == b'\'') {
61                current.push('\'');
62                index += 1;
63                while index < bytes.len() {
64                    let b = bytes[index];
65                    current.push(b as char);
66                    index += 1;
67                    if b == b'\'' {
68                        // Doubled quote is an escaped quote; consume the next
69                        // quote too if it's a quote.
70                        if index < bytes.len() && bytes[index] == b'\'' {
71                            current.push('\'');
72                            index += 1;
73                            continue;
74                        }
75                        break;
76                    }
77                }
78                continue;
79            }
80            // Double-quoted identifier: `"..."` with `""` escapes.
81            if bytes_rest.first().is_some_and(|&b| b == b'"') {
82                current.push('"');
83                index += 1;
84                while index < bytes.len() {
85                    let b = bytes[index];
86                    current.push(b as char);
87                    index += 1;
88                    if b == b'"' {
89                        if index < bytes.len() && bytes[index] == b'"' {
90                            current.push('"');
91                            index += 1;
92                            continue;
93                        }
94                        break;
95                    }
96                }
97                continue;
98            }
99            // Top-level statement separator.
100            if bytes_rest.first().is_some_and(|&b| b == b';') {
101                let trimmed = current.trim();
102                if !trimmed.is_empty() {
103                    statements.push(Self {
104                        sql: trimmed.to_owned(),
105                    });
106                }
107                current.clear();
108                index += 1;
109                continue;
110            }
111            // Normal byte: copy through.
112            let ch = bytes[index] as char;
113            current.push(ch);
114            index += 1;
115        }
116        let trimmed = current.trim();
117        if !trimmed.is_empty() {
118            statements.push(Self {
119                sql: trimmed.to_owned(),
120            });
121        }
122        statements
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::SqlStatement;
129
130    #[test]
131    fn splits_simple_statements() {
132        let stmts = SqlStatement::split("SELECT 1; SELECT 2;");
133        assert_eq!(stmts.len(), 2);
134        assert_eq!(stmts[0].sql, "SELECT 1");
135        assert_eq!(stmts[1].sql, "SELECT 2");
136    }
137
138    #[test]
139    fn strips_line_and_block_comments() {
140        let stmts =
141            SqlStatement::split("-- comment\nSELECT 1; /* block */ SELECT 2 -- trailing\n;");
142        assert_eq!(stmts.len(), 2);
143        assert_eq!(stmts[0].sql, "SELECT 1");
144        assert_eq!(stmts[1].sql, "SELECT 2");
145    }
146
147    #[test]
148    fn does_not_split_on_semicolon_in_string() {
149        let stmts = SqlStatement::split("INSERT INTO t VALUES ('a;b');");
150        assert_eq!(stmts.len(), 1);
151        assert_eq!(stmts[0].sql, "INSERT INTO t VALUES ('a;b')");
152    }
153
154    #[test]
155    fn handles_escaped_quotes_in_string() {
156        let stmts = SqlStatement::split("INSERT INTO t VALUES ('it''s ok');");
157        assert_eq!(stmts.len(), 1);
158        assert!(stmts[0].sql.contains("it''s ok"));
159    }
160
161    #[test]
162    fn handles_unterminated_string_without_panicking() {
163        let stmts = SqlStatement::split("INSERT INTO t VALUES ('unterminated");
164        assert_eq!(stmts.len(), 1);
165        assert!(stmts[0].sql.contains("unterminated"));
166    }
167
168    #[test]
169    fn handles_quoted_identifiers() {
170        let stmts = SqlStatement::split("ALTER TABLE \"my;table\" ADD COLUMN x int;");
171        assert_eq!(stmts.len(), 1);
172    }
173
174    #[test]
175    fn empty_input_yields_no_statements() {
176        assert!(SqlStatement::split("").is_empty());
177        assert!(SqlStatement::split("-- only a comment").is_empty());
178        assert!(SqlStatement::split("   ;  ;  ").is_empty());
179    }
180}