SQL Playground — Python source

"""sql-playground — polyglot showcase port (Python).

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.

Display source — part of CosmoDev's polyglot tool pages.

Type notes
----------
Python distinguishes ``int`` and ``float`` natively, and treats ``True``/
``False``/``None`` as singletons. Cells may be any of those plus ``str``;
anything else (list/dict/object) is JSON-encoded for display, mirroring the
TypeScript behaviour where any non-primitive becomes JSON.
"""

from __future__ import annotations

import json
import re
from typing import Any, Iterable, List, Optional, Sequence, Tuple

__all__ = [
    "SchemaTable",
    "split_statements",
    "is_read_only_statement",
    "format_scalar",
    "rows_to_csv",
    "rows_to_markdown",
    "rows_to_json",
    "summarize_schema",
]


class SchemaTable:
    """One extracted CREATE TABLE summary: its name and lower-cased columns.

    Implemented as a small dataclass-like value object so callers can use
    ``schema[i].table`` / ``schema[i].columns`` the same way they would access
    the ``{ table, columns }`` object in the TypeScript original.
    """

    __slots__ = ("table", "columns")

    def __init__(self, table: str, columns: Sequence[str]) -> None:
        self.table = table
        self.columns = list(columns)

    # Equality/repr keep parity with TS structural object comparison handy for
    # any future cross-language test harness.
    def __eq__(self, other: object) -> bool:
        return (
            isinstance(other, SchemaTable)
            and self.table == other.table
            and self.columns == other.columns
        )

    def __repr__(self) -> str:  # pragma: no cover - debug aid only
        return f"SchemaTable(table={self.table!r}, columns={self.columns!r})"

    def to_dict(self) -> dict:
        """JSON-friendly dict form (matches the TS ``{ table, columns }``)."""
        return {"table": self.table, "columns": self.columns}


# ---------------------------------------------------------------------------
# Statement splitting
# ---------------------------------------------------------------------------


def split_statements(sql: str) -> List[str]:
    """Split a SQL script into top-level statements on ``;``.

    Single-quote string literals are respected: a doubled ``''`` inside a
    string is an escaped literal quote (SQL standard) and does not end the
    string. Whitespace-only statements are dropped.
    """
    stmts: List[str] = []
    current = ""
    in_string = False
    n = len(sql)

    i = 0
    while i < n:
        ch = sql[i]

        if in_string:
            current += ch
            if ch == "'":
                if i + 1 < n and sql[i + 1] == "'":
                    # Escaped literal quote: consume both chars.
                    current += sql[i + 1]
                    i += 1
                else:
                    in_string = False
            i += 1
            continue

        if ch == "'":
            in_string = True
            current += ch
        elif ch == ";":
            trimmed = current.strip()
            if trimmed:
                stmts.append(trimmed)
            current = ""
        else:
            current += ch
        i += 1

    # Flush any trailing statement that wasn't terminated by ";".
    trimmed = current.strip()
    if trimmed:
        stmts.append(trimmed)

    return stmts


# ---------------------------------------------------------------------------
# Read-only classification
# ---------------------------------------------------------------------------

# `\b` after the alternation enforces a word boundary, so "SELECTOR" doesn't
# match "SELECT". The `re.I` flag makes it case-insensitive.
_READONLY_RE = re.compile(r"^(SELECT|WITH|VALUES|EXPLAIN|PRAGMA)\b", re.IGNORECASE)


def is_read_only_statement(sql: str) -> bool:
    """True when the statement is read-only (safe on a snapshot).

    Recognises the common query-leading keywords: SELECT, WITH, VALUES,
    EXPLAIN, PRAGMA.
    """
    return _READONLY_RE.match(sql.strip()) is not None


# ---------------------------------------------------------------------------
# Scalar formatting
# ---------------------------------------------------------------------------


def format_scalar(v: Any) -> str:
    """Render one cell value as the textual form used in tables/CSV.

    - ``None``            -> ``"NULL"``
    - ``bool``            -> ``"True"``/``"False"`` (Python's native str)
    - ``int``/``float``   -> native string form
    - ``str``             -> verbatim
    - anything else       -> JSON

    Note on bool: Python's ``str(True)`` is ``"True"``/``"False"`` (capitalised),
    whereas JavaScript's is ``"true"``/``"false"``. We intentionally keep
    Python's native form so the snippet reads idiomatically; if exact byte
    parity with the TS lib is required, downcast the bool first.
    """
    if v is None:
        return "NULL"
    # Order matters: bool is a subclass of int in Python, so check it first.
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, (int, float)):
        return _number_to_str(v)
    if isinstance(v, str):
        return v
    # Arrays / objects -> JSON. ``default=str`` is a safety net for values that
    # aren't natively JSON-serialisable (e.g. datetime, Decimal); the TS path
    # would also fall back to a stringified form via JSON.stringify replacers.
    return json.dumps(v, default=str)


def _number_to_str(v: "int | float") -> str:
    """Format a number the way JavaScript's ``String(number)`` would.

    Whole-valued floats render without a fractional tail (e.g. ``42`` not
    ``42.0``), matching the TS lib's ``String(v)`` for the values that show up
    in SQL result sets. ``float('inf')`` / ``nan`` become ``Infinity``/``NaN``
    like JS, not Python's ``"inf"``.
    """
    if isinstance(v, bool):  # pragma: no cover - guarded above
        return "true" if v else "false"
    if isinstance(v, int):
        return str(v)
    # float
    if v != v:  # NaN
        return "NaN"
    if v == float("inf"):
        return "Infinity"
    if v == float("-inf"):
        return "-Infinity"
    if v.is_integer() and abs(v) < 1e21:
        return str(int(v))
    # repr() yields the shortest round-trip representation in Python 3.
    return repr(v)


# ---------------------------------------------------------------------------
# CSV
# ---------------------------------------------------------------------------

# A field needs CSV quoting if it contains a comma, double quote, or any
# newline flavour (CR / LF / CRLF).
_CSV_QUOTE_RE = re.compile(r'[,"\n\r]')


def _csv_field(v: Any) -> str:
    """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.
    """
    s = format_scalar(v)
    if _CSV_QUOTE_RE.search(s):
        return '"' + s.replace('"', '""') + '"'
    return s


def rows_to_csv(columns: Sequence[str], rows: Iterable[Sequence[Any]]) -> str:
    """Serialize a result set to CSV (header + one row per record, trailing LF)."""
    lines = [",".join(_csv_field(c) for c in columns)]
    for row in rows:
        lines.append(",".join(_csv_field(cell) for cell in row))
    return "\n".join(lines) + "\n"


# ---------------------------------------------------------------------------
# Markdown
# ---------------------------------------------------------------------------


def _md_escape(s: str) -> str:
    """Escape a pipe so it doesn't break the Markdown table layout."""
    return s.replace("|", r"\|")


def rows_to_markdown(columns: Sequence[str], rows: Iterable[Sequence[Any]]) -> str:
    """Serialize a result set as a GitHub-flavored Markdown table."""
    header = "| " + " | ".join(_md_escape(c) for c in columns) + " |"
    sep = "| " + " | ".join("---" for _ in columns) + " |"

    lines = [header, sep]
    for row in rows:
        cells = " | ".join(_md_escape(format_scalar(cell)) for cell in row)
        lines.append("| " + cells + " |")
    return "\n".join(lines) + "\n"


# ---------------------------------------------------------------------------
# JSON
# ---------------------------------------------------------------------------


def rows_to_json(columns: Sequence[str], rows: Iterable[Sequence[Any]]) -> str:
    """Serialize a result set as a pretty-printed JSON array of objects.

    Objects are keyed by column name. Missing cells (row shorter than columns)
    become ``null``, mirroring the TS ``row[i] ?? null``.
    """
    out = []
    for row in rows:
        obj = {}
        for i, col in enumerate(columns):
            # ``row[i] if i < len(row) else None`` mirrors the JS nullish
            # coalescing; an explicit None in the row also serialises as null.
            obj[col] = row[i] if i < len(row) else None
        out.append(obj)
    # indent=2 produces the same 2-space layout as JSON.stringify(_, null, 2).
    return json.dumps(out, indent=2, default=str)


# ---------------------------------------------------------------------------
# Schema summary
# ---------------------------------------------------------------------------

# Matches a table-level constraint keyword at the start of a column definition
# line (PRIMARY KEY, FOREIGN KEY, UNIQUE, CHECK, CONSTRAINT) -> not a column.
_CONSTRAINT_RE = re.compile(
    r"^(PRIMARY\s+KEY|FOREIGN\s+KEY|UNIQUE|CHECK|CONSTRAINT)\b", re.IGNORECASE
)

# Captures the leading identifier of a column definition, with optional
# surrounding quote/backtick. ``\w`` covers [A-Za-z0-9_].
_COLUMN_NAME_RE = re.compile(r"""^["'`]?(\w+)["'`]?""")

# Finds each "CREATE TABLE [IF NOT EXISTS] <name> (" header. Used with
# finditer + match.end() to walk the script while tracking each match's span.
_HEADER_RE = re.compile(
    r"""CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?["'`]?(\w+)["'`]?\s*\(""",
    re.IGNORECASE,
)


def _extract_column_name(defn: str) -> Optional[str]:
    """Pull the column name (lower-cased) from one column definition.

    Returns ``None`` for table-level constraint lines, which are not columns.
    """
    defn = defn.strip()
    if not defn:
        return None
    if _CONSTRAINT_RE.match(defn):
        return None
    m = _COLUMN_NAME_RE.match(defn)
    if not m:
        return None
    return m.group(1).lower()


def summarize_schema(create_sql: str) -> List[SchemaTable]:
    """Best-effort extraction of ``SchemaTable`` 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 list on any error.
    """
    try:
        results: List[SchemaTable] = []
        n = len(create_sql)

        for header in _HEADER_RE.finditer(create_sql):
            table = header.group(1).lower()
            body_start = header.end()

            # Walk forward to find the matching close paren of the table body.
            depth = 1
            i = body_start
            while i < n and depth > 0:
                c = create_sql[i]
                if c == "(":
                    depth += 1
                elif c == ")":
                    depth -= 1
                i += 1
            if depth != 0:
                # Unbalanced parens -> malformed; skip this match.
                continue

            body = create_sql[body_start : i - 1]
            columns: List[str] = []

            # Split body on top-level commas only (depth-tracked) so that
            # commas inside ``NUMERIC(10, 2)`` or ``CHECK(x IN (1,2))`` stay
            # within their column definition.
            col_depth = 0
            current_chars: List[str] = []
            for ch in body:
                if ch == "(":
                    col_depth += 1
                    current_chars.append(ch)
                elif ch == ")":
                    col_depth -= 1
                    current_chars.append(ch)
                elif ch == "," and col_depth == 0:
                    col = _extract_column_name("".join(current_chars))
                    if col:
                        columns.append(col)
                    current_chars = []
                else:
                    current_chars.append(ch)

            # Flush the trailing column definition.
            last_col = _extract_column_name("".join(current_chars))
            if last_col:
                columns.append(last_col)

            results.append(SchemaTable(table, columns))

        return results
    except Exception:
        # The TS original swallows any thrown error and returns []; mirror
        # that contract so a malformed script never crashes the tool.
        return []