SQL Playground — Rust source

// sql-playground — polyglot showcase port (Rust)
//
// Pure helper logic for the CosmoDev "SQL Playground" tool: splitting a SQL
// script into statements, classifying read-only ones, serializing result sets
// to CSV / Markdown / JSON, and extracting a schema summary from CREATE TABLE
// scripts.
//
// Ported from src/lib/sql-playground.ts (the canonical TypeScript lib that
// powers the live tool). Functionally equivalent: same inputs -> same outputs.
//
// This showcase uses the standard library only. JSON encoding is hand-rolled
// (no serde) so the snippet stays self-contained and dependency-free.
//
// Display source — part of CosmoDev's polyglot tool pages.

use std::fmt::Write as _;

/// One extracted CREATE TABLE summary: its name and the lower-cased column
/// names declared in its body.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SchemaTable {
    pub table: String,
    pub columns: Vec<String>,
}

/// A scalar value held in a result-set cell. Mirrors the small set of types
/// the TypeScript version handles: null, bool, number, string, or anything
/// else (serialised as JSON).
#[derive(Debug, Clone)]
pub enum Scalar {
    Null,
    Bool(bool),
    Int(i64),
    Float(f64),
    Str(String),
    /// Already-JSON-encoded representation of a non-scalar (object/array).
    Json(String),
}

impl Scalar {
    /// Build a Scalar from a Rust value, mapping the common primitives.
    /// Anything that doesn't fit one of the typed variants is JSON-encoded.
    pub fn from_any<T: ScalarFrom>(v: T) -> Self {
        v.into_scalar()
    }
}

/// Helper trait so `Scalar::from_any` accepts several native types.
pub trait ScalarFrom {
    fn into_scalar(self) -> Scalar;
}

impl ScalarFrom for () {
    fn into_scalar(self) -> Scalar {
        Scalar::Null
    }
}
impl ScalarFrom for bool {
    fn into_scalar(self) -> Scalar {
        Scalar::Bool(self)
    }
}
impl ScalarFrom for i64 {
    fn into_scalar(self) -> Scalar {
        Scalar::Int(self)
    }
}
impl ScalarFrom for f64 {
    fn into_scalar(self) -> Scalar {
        Scalar::Float(self)
    }
}
impl ScalarFrom for &str {
    fn into_scalar(self) -> Scalar {
        Scalar::Str(self.to_string())
    }
}
impl ScalarFrom for String {
    fn into_scalar(self) -> Scalar {
        Scalar::Str(self)
    }
}
impl<T: ScalarFrom> ScalarFrom for Option<T> {
    fn into_scalar(self) -> Scalar {
        match self {
            Some(v) => v.into_scalar(),
            None => Scalar::Null,
        }
    }
}

/// Render one cell as the textual form used in tables/CSV:
/// null -> "NULL", bool/int/float -> native string, string -> verbatim,
/// anything else -> its JSON form.
pub fn format_scalar(v: &Scalar) -> String {
    match v {
        Scalar::Null => "NULL".to_string(),
        Scalar::Bool(b) => b.to_string(),
        Scalar::Int(i) => i.to_string(),
        // Rust's default float formatter already produces shortest round-trip
        // representations (e.g. `3.14`, `1` for whole floats), matching JS's
        // Number -> String for the values that typically appear in result sets.
        Scalar::Float(f) => format_float_like_js(*f),
        Scalar::Str(s) => s.clone(),
        Scalar::Json(j) => j.clone(),
    }
}

/// Format an f64 the way JavaScript's `String(number)` would for the values
/// that show up in SQL result sets: whole numbers render without a fractional
/// tail (e.g. `42` not `42.0`), and f64's default Display gives that.
fn format_float_like_js(f: f64) -> String {
    if f.is_nan() {
        return "NaN".to_string();
    }
    if f.is_infinite() {
        return if f > 0.0 { "Infinity".to_string() } else { "-Infinity".to_string() };
    }
    if f == f.trunc() && f.abs() < 1e21 {
        // Whole number -> integer form, no trailing ".0".
        format!("{}", f as i64)
    } else {
        format!("{}", f)
    }
}

// ---------------------------------------------------------------------------
// Statement splitting
// ---------------------------------------------------------------------------

/// Split a SQL script into top-level statements on `;`, respecting single-quote
/// string literals. A doubled `''` inside a string is an escaped literal quote
/// (SQL standard) and does not end the string. Whitespace-only statements are
/// dropped.
pub fn split_statements(sql: &str) -> Vec<String> {
    let chars: Vec<char> = sql.chars().collect();
    let mut stmts = Vec::new();
    let mut current = String::new();
    let mut in_string = false;

    let mut i = 0;
    while i < chars.len() {
        let ch = chars[i];

        if in_string {
            current.push(ch);
            if ch == '\'' {
                if i + 1 < chars.len() && chars[i + 1] == '\'' {
                    // Escaped literal quote: consume both chars.
                    current.push(chars[i + 1]);
                    i += 1;
                } else {
                    in_string = false;
                }
            }
        } else if ch == '\'' {
            in_string = true;
            current.push(ch);
        } else if ch == ';' {
            let trimmed = current.trim();
            if !trimmed.is_empty() {
                stmts.push(trimmed.to_string());
            }
            current.clear();
        } else {
            current.push(ch);
        }
        i += 1;
    }

    // Flush any trailing statement that wasn't terminated by `;`.
    let trimmed = current.trim();
    if !trimmed.is_empty() {
        stmts.push(trimmed.to_string());
    }

    stmts
}

// ---------------------------------------------------------------------------
// Read-only classification
// ---------------------------------------------------------------------------

/// True when the statement is read-only (safe to run against a snapshot
/// without mutating state). Recognises the common query-leading keywords:
/// SELECT, WITH, VALUES, EXPLAIN, PRAGMA. Case-insensitive, and requires a
/// word boundary after the keyword so e.g. "SELECTOR" doesn't match "SELECT".
pub fn is_read_only_statement(sql: &str) -> bool {
    let upper = sql.trim_start().to_uppercase();
    for kw in ["SELECT", "WITH", "VALUES", "EXPLAIN", "PRAGMA"] {
        let rest = match upper.strip_prefix(kw) {
            Some(r) => r,
            None => continue,
        };
        match rest.chars().next() {
            // End of input, or a non-word char -> word boundary satisfied.
            None => return true,
            Some(c) if !c.is_alphanumeric() && c != '_' => return true,
            _ => {}
        }
    }
    false
}

// ---------------------------------------------------------------------------
// CSV
// ---------------------------------------------------------------------------

/// Quote one cell per RFC-4180-ish CSV: if the formatted value contains a
/// comma, double quote, or newline, wrap it in double quotes and double any
/// embedded double quotes.
fn csv_field(v: &Scalar) -> String {
    let s = format_scalar(v);
    if s.chars().any(|c| matches!(c, ',' | '"' | '\n' | '\r')) {
        let escaped = s.replace('"', "\"\"");
        format!("\"{}\"", escaped)
    } else {
        s
    }
}

/// Serialize a result set to CSV (header row + one row per record, trailing LF).
pub fn rows_to_csv(columns: &[String], rows: &[Vec<Scalar>]) -> String {
    let mut out = String::new();
    let header = columns
        .iter()
        .map(|c| csv_field(&Scalar::Str(c.clone())))
        .collect::<Vec<_>>()
        .join(",");
    out.push_str(&header);
    out.push('\n');
    for row in rows {
        let line = row.iter().map(csv_field).collect::<Vec<_>>().join(",");
        out.push_str(&line);
        out.push('\n');
    }
    out
}

// ---------------------------------------------------------------------------
// Markdown
// ---------------------------------------------------------------------------

/// Escape a pipe so it doesn't break the Markdown table layout.
fn md_escape(s: &str) -> String {
    s.replace('|', "\\|")
}

/// Serialize a result set as a GitHub-flavored Markdown table.
pub fn rows_to_markdown(columns: &[String], rows: &[Vec<Scalar>]) -> String {
    let header_cells: Vec<String> = columns.iter().map(|c| md_escape(c)).collect();
    let header = format!("| {} |", header_cells.join(" | "));
    let sep_cells: Vec<String> = columns.iter().map(|_| "---".to_string()).collect();
    let sep = format!("| {} |", sep_cells.join(" | "));

    let mut out = String::new();
    writeln!(out, "{}", header).unwrap();
    writeln!(out, "{}", sep).unwrap();
    for row in rows {
        let cells: Vec<String> = row
            .iter()
            .map(|c| md_escape(&format_scalar(c)))
            .collect();
        writeln!(out, "| {} |", cells.join(" | ")).unwrap();
    }
    out
}

// ---------------------------------------------------------------------------
// JSON
// ---------------------------------------------------------------------------

/// JSON-encode a string literal body (without surrounding quotes added here).
/// Escapes the standard control chars and the double quote/backslash.
fn json_escape_into(out: &mut String, s: &str) {
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            '\x08' => out.push_str("\\b"),
            '\x0c' => out.push_str("\\f"),
            c if (c as u32) < 0x20 => {
                write!(out, "\\u{:04x}", c as u32).unwrap();
            }
            c => out.push(c),
        }
    }
}

/// Render a single Scalar as its JSON value text (no surrounding container).
fn scalar_to_json(v: &Scalar) -> String {
    match v {
        Scalar::Null => "null".to_string(),
        Scalar::Bool(b) => b.to_string(),
        Scalar::Int(i) => i.to_string(),
        Scalar::Float(f) => {
            if f.is_nan() || f.is_infinite() {
                "null".to_string()
            } else if *f == f.trunc() && f.abs() < 1e21 {
                format!("{}", *f as i64)
            } else {
                format!("{}", f)
            }
        }
        Scalar::Str(s) => {
            let mut buf = String::from("\"");
            json_escape_into(&mut buf, s);
            buf.push('"');
            buf
        }
        Scalar::Json(j) => j.clone(),
    }
}

/// Serialize a result set as a pretty-printed (2-space) JSON array of objects
/// keyed by column name. Missing cells (row shorter than columns) become null.
pub fn rows_to_json(columns: &[String], rows: &[Vec<Scalar>]) -> String {
    let mut out = String::new();
    if rows.is_empty() {
        out.push_str("[]");
        return out;
    }
    out.push_str("[\n");
    for (ri, row) in rows.iter().enumerate() {
        out.push_str("  {\n");
        for (ci, col) in columns.iter().enumerate() {
            out.push_str("    \"");
            json_escape_into(&mut out, col);
            out.push_str("\": ");
            let cell = row.get(ci).unwrap_or(&Scalar::Null);
            out.push_str(&scalar_to_json(cell));
            if ci + 1 < columns.len() {
                out.push(',');
            }
            out.push('\n');
        }
        out.push_str("  }");
        if ri + 1 < rows.len() {
            out.push(',');
        }
        out.push('\n');
    }
    out.push(']');
    out
}

// ---------------------------------------------------------------------------
// Schema summary
// ---------------------------------------------------------------------------

/// Pull the column name (lower-cased) from a single CREATE TABLE column
/// definition. Returns None for table-level constraint lines (PRIMARY KEY,
/// FOREIGN KEY, UNIQUE, CHECK, CONSTRAINT), which are not columns.
fn extract_column_name(def: &str) -> Option<String> {
    let def = def.trim();
    if def.is_empty() {
        return None;
    }
    if is_constraint_line(def) {
        return None;
    }
    // Match an optional opening quote/backtick, then \w+, then an optional
    // closing quote/backtick — capturing just the identifier.
    let mut chars = def.chars();
    let first = chars.next()?;
    let name: String = if first == '"' || first == '\'' || first == '`' {
        chars.take_while(|c| *c != first).collect()
    } else {
        // \w is [A-Za-z0-9_]; take leading word chars and prepend the first.
        let mut s = String::new();
        s.push(first);
        for c in chars {
            if c.is_alphanumeric() || c == '_' {
                s.push(c);
            } else {
                break;
            }
        }
        s
    };
    if name.is_empty() || name.chars().any(|c| !(c.is_alphanumeric() || c == '_')) {
        return None;
    }
    Some(name.to_lowercase())
}

/// Case-insensitive prefix check for table-level constraint keywords.
fn is_constraint_line(def: &str) -> bool {
    let upper = def.to_uppercase();
    let collapsed = upper.replace('\n', " ");
    let collapsed = collapsed.split_whitespace().collect::<Vec<_>>().join(" ");
    for kw in [
        "PRIMARY KEY",
        "FOREIGN KEY",
        "UNIQUE",
        "CHECK",
        "CONSTRAINT",
    ] {
        if collapsed.starts_with(kw) {
            return true;
        }
    }
    false
}

/// Find the next `CREATE TABLE name (` header on or after `from`, returning
/// (table_name_lowercased, body_open_paren_index) if found.
fn find_next_header(sql: &str, from: usize) -> Option<(String, usize)> {
    let bytes = sql.as_bytes();
    // Lowercased copy for case-insensitive scanning; we still index into the
    // original (both have identical byte lengths for ASCII SQL — and CREATE
    // TABLE keywords are ASCII by spec).
    let lower: String = sql.chars().map(|c| c.to_ascii_uppercase()).collect();
    let lower_bytes = lower.as_bytes();

    let mut i = from;
    while i + 12 <= lower_bytes.len() {
        // Look for "CREATE" then whitespace "TABLE".
        if lower_bytes[i..].starts_with(b"CREATE") {
            let after = i + 6;
            if let Some(ws_end) = next_non_ws(lower_bytes, after) {
                if lower_bytes[ws_end..].starts_with(b"TABLE") {
                    let after_table = ws_end + 5;
                    if let Some(rest_start) = next_non_ws(lower_bytes, after_table) {
                        let (name, after_name) = read_identifier(bytes, rest_start)?;
                        let after_name_ws = next_non_ws(lower_bytes, after_name)?;
                        if bytes.get(after_name_ws) == Some(&b'(') {
                            let body_start = after_name_ws + 1;
                            return Some((name.to_lowercase(), body_start));
                        }
                    }
                }
            }
        }
        i += 1;
    }
    None
}

fn next_non_ws(bytes: &[u8], mut i: usize) -> Option<usize> {
    while i < bytes.len() && (bytes[i] == b' ' || bytes[i] == b'\t' || bytes[i] == b'\n' || bytes[i] == b'\r') {
        i += 1;
    }
    if i >= bytes.len() {
        None
    } else {
        Some(i)
    }
}

/// Read an `\w+` identifier, optionally wrapped in matching quotes/backticks,
/// handling the `IF NOT EXISTS` clause between the table name and `(`. Returns
/// the name and the index just past the identifier.
fn read_identifier(bytes: &[u8], start: usize) -> Option<(&str, usize)> {
    if start >= bytes.len() {
        return None;
    }
    let first = bytes[start];
    let (name, end) = if first == b'"' || first == b'\'' || first == b'`' {
        // Quoted identifier: up to the matching close.
        let mut j = start + 1;
        while j < bytes.len() && bytes[j] != first {
            j += 1;
        }
        if j >= bytes.len() {
            return None;
        }
        let name = std::str::from_utf8(&bytes[start + 1..j]).ok()?;
        (name, j + 1)
    } else {
        // Bare \w+ identifier.
        let mut j = start;
        while j < bytes.len()
            && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_')
        {
            j += 1;
        }
        if j == start {
            return None;
        }
        let name = std::str::from_utf8(&bytes[start..j]).ok()?;
        (name, j)
    };
    // Skip optional "IF NOT EXISTS" between name and "(".
    let mut k = end;
    if let Some(ws) = next_non_ws(bytes, k) {
        let tail = &bytes[ws..];
        let upper: String = tail
            .iter()
            .take(12)
            .map(|&b| b.to_ascii_uppercase() as char)
            .collect();
        if upper.starts_with("IF NOT EXISTS") || upper.starts_with("IF NOT EXIS") {
            // Skip "IF NOT EXISTS".
            k = ws + "IF NOT EXISTS".len();
        }
    }
    Some((name, k))
}

/// Best-effort extraction of `{ table, columns }` summaries from a CREATE TABLE
/// script. Handles IF NOT EXISTS, quoted identifiers, and parenthesised
/// type/constraint bodies (depth-tracked so a `,` inside `NUMERIC(10,2)` doesn't
/// split a column). Returns an empty slice on any error.
pub fn summarize_schema(create_sql: &str) -> Vec<SchemaTable> {
    let mut results = Vec::new();
    let bytes = create_sql.as_bytes();
    let mut cursor = 0;

    while cursor < bytes.len() {
        let (table, body_start) = match find_next_header(create_sql, cursor) {
            Some(h) => h,
            None => break,
        };

        // Walk forward to find the matching close paren of the table body.
        let mut depth: i32 = 1;
        let mut i = body_start;
        while i < bytes.len() && depth > 0 {
            match bytes[i] {
                b'(' => depth += 1,
                b')' => depth -= 1,
                _ => {}
            }
            i += 1;
        }
        if depth != 0 {
            // Unbalanced parens -> malformed; skip and resume scanning.
            cursor = body_start;
            continue;
        }

        let body_end = i - 1; // index of the closing ')'
        let body = &create_sql[body_start..body_end];

        // Split the body on top-level commas only (depth-tracked).
        let mut columns: Vec<String> = Vec::new();
        let mut current = String::new();
        let mut col_depth: i32 = 0;
        for ch in body.chars() {
            match ch {
                '(' => {
                    col_depth += 1;
                    current.push(ch);
                }
                ')' => {
                    col_depth -= 1;
                    current.push(ch);
                }
                ',' if col_depth == 0 => {
                    if let Some(col) = extract_column_name(&current) {
                        columns.push(col);
                    }
                    current.clear();
                }
                _ => current.push(ch),
            }
        }
        if let Some(col) = extract_column_name(&current) {
            columns.push(col);
        }

        results.push(SchemaTable { table, columns });
        cursor = i;
    }

    results
}