arcature-data 2026.2.0

//! SQL statement tokenization: split a script into [`SqlStatement`]s,
//! respecting string literals, quoted identifiers, and comments.
//!
//! The splitter is a **robust, conservative scanner** — not a full SQL
//! parser. It walks the input byte by byte, skipping line comments (`--`)
//! and block comments (`/* */`), tracking single-quoted string literals
//! (with `''` escapes) and double-quoted identifiers, and splitting on
//! top-level `;`. It never panics on hostile or malformed input: an
//! unterminated string or comment simply ends at end-of-input, and the
//! accumulated text becomes the final statement.

/// A single SQL statement, normalized for classification.
///
/// Construct via [`SqlStatement::split`] (a script into statements) or
/// [`SqlStatement::single`] (one statement). The `sql` field carries the
/// normalized text the classifier scans; comments are stripped, but quoted
/// literals and identifiers are preserved so keyword detection does not fire
/// inside a string.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SqlStatement {
    /// The normalized statement text (comments stripped; statements trimmed).
    pub sql: String,
}

impl SqlStatement {
    /// Wrap a single already-split statement.
    #[must_use]
    pub fn single(sql: impl Into<String>) -> Self {
        let sql = sql.into();
        Self {
            sql: sql.trim().to_owned(),
        }
    }

    /// Split a SQL script into statements. Whitespace-only statements are
    /// dropped. Never panics.
    #[must_use]
    pub fn split(script: &str) -> Vec<Self> {
        let mut statements = Vec::new();
        let mut current = String::new();
        let bytes = script.as_bytes();
        let mut index = 0;
        while index < bytes.len() {
            let rest = &script[index..];
            // Line comment: `--` to end of line.
            if rest.starts_with("--") {
                let end = rest.find('\n').map(|n| n + 1).unwrap_or(rest.len());
                index += end;
                continue;
            }
            // Block comment: `/* ... */` (nesting not supported by PostgreSQL
            // block comments, so a simple non-nesting scan is correct).
            if rest.starts_with("/*") {
                let end = rest.find("*/").map(|n| n + 2).unwrap_or(rest.len());
                index += end;
                continue;
            }
            // Single-quoted string literal: `'...'` with `''` escapes.
            let bytes_rest = &bytes[index..];
            if bytes_rest.first().is_some_and(|&b| b == b'\'') {
                current.push('\'');
                index += 1;
                while index < bytes.len() {
                    let b = bytes[index];
                    current.push(b as char);
                    index += 1;
                    if b == b'\'' {
                        // Doubled quote is an escaped quote; consume the next
                        // quote too if it's a quote.
                        if index < bytes.len() && bytes[index] == b'\'' {
                            current.push('\'');
                            index += 1;
                            continue;
                        }
                        break;
                    }
                }
                continue;
            }
            // Double-quoted identifier: `"..."` with `""` escapes.
            if bytes_rest.first().is_some_and(|&b| b == b'"') {
                current.push('"');
                index += 1;
                while index < bytes.len() {
                    let b = bytes[index];
                    current.push(b as char);
                    index += 1;
                    if b == b'"' {
                        if index < bytes.len() && bytes[index] == b'"' {
                            current.push('"');
                            index += 1;
                            continue;
                        }
                        break;
                    }
                }
                continue;
            }
            // Top-level statement separator.
            if bytes_rest.first().is_some_and(|&b| b == b';') {
                let trimmed = current.trim();
                if !trimmed.is_empty() {
                    statements.push(Self {
                        sql: trimmed.to_owned(),
                    });
                }
                current.clear();
                index += 1;
                continue;
            }
            // Normal byte: copy through.
            let ch = bytes[index] as char;
            current.push(ch);
            index += 1;
        }
        let trimmed = current.trim();
        if !trimmed.is_empty() {
            statements.push(Self {
                sql: trimmed.to_owned(),
            });
        }
        statements
    }
}

#[cfg(test)]
mod tests {
    use super::SqlStatement;

    #[test]
    fn splits_simple_statements() {
        let stmts = SqlStatement::split("SELECT 1; SELECT 2;");
        assert_eq!(stmts.len(), 2);
        assert_eq!(stmts[0].sql, "SELECT 1");
        assert_eq!(stmts[1].sql, "SELECT 2");
    }

    #[test]
    fn strips_line_and_block_comments() {
        let stmts =
            SqlStatement::split("-- comment\nSELECT 1; /* block */ SELECT 2 -- trailing\n;");
        assert_eq!(stmts.len(), 2);
        assert_eq!(stmts[0].sql, "SELECT 1");
        assert_eq!(stmts[1].sql, "SELECT 2");
    }

    #[test]
    fn does_not_split_on_semicolon_in_string() {
        let stmts = SqlStatement::split("INSERT INTO t VALUES ('a;b');");
        assert_eq!(stmts.len(), 1);
        assert_eq!(stmts[0].sql, "INSERT INTO t VALUES ('a;b')");
    }

    #[test]
    fn handles_escaped_quotes_in_string() {
        let stmts = SqlStatement::split("INSERT INTO t VALUES ('it''s ok');");
        assert_eq!(stmts.len(), 1);
        assert!(stmts[0].sql.contains("it''s ok"));
    }

    #[test]
    fn handles_unterminated_string_without_panicking() {
        let stmts = SqlStatement::split("INSERT INTO t VALUES ('unterminated");
        assert_eq!(stmts.len(), 1);
        assert!(stmts[0].sql.contains("unterminated"));
    }

    #[test]
    fn handles_quoted_identifiers() {
        let stmts = SqlStatement::split("ALTER TABLE \"my;table\" ADD COLUMN x int;");
        assert_eq!(stmts.len(), 1);
    }

    #[test]
    fn empty_input_yields_no_statements() {
        assert!(SqlStatement::split("").is_empty());
        assert!(SqlStatement::split("-- only a comment").is_empty());
        assert!(SqlStatement::split("   ;  ;  ").is_empty());
    }
}