SQL Playground — Go source

// sql-playground — polyglot showcase port (Go)
//
// 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.
package sqlplayground

import (
	"encoding/json"
	"regexp"
	"strconv"
	"strings"
)

// SchemaTable is the { table, columns } summary extracted from a CREATE TABLE
// statement.
type SchemaTable struct {
	Table   string   `json:"table"`
	Columns []string `json:"columns"`
}

// SplitStatements splits 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.
func SplitStatements(sql string) []string {
	var stmts []string
	var current strings.Builder
	inString := false
	runes := []rune(sql)

	for i := 0; i < len(runes); i++ {
		ch := runes[i]

		switch {
		case inString:
			current.WriteRune(ch)
			if ch == '\'' {
				if i+1 < len(runes) && runes[i+1] == '\'' {
					// Escaped literal quote: consume both runes.
					current.WriteRune(runes[i+1])
					i++
				} else {
					inString = false
				}
			}
		case ch == '\'':
			inString = true
			current.WriteRune(ch)
		case ch == ';':
			trimmed := strings.TrimSpace(current.String())
			if trimmed != "" {
				stmts = append(stmts, trimmed)
			}
			current.Reset()
		default:
			current.WriteRune(ch)
		}
	}

	// Flush any trailing statement that wasn't terminated by ';'.
	trimmed := strings.TrimSpace(current.String())
	if trimmed != "" {
		stmts = append(stmts, trimmed)
	}

	// Guarantee a non-nil slice for parity with empty arrays elsewhere.
	if stmts == nil {
		stmts = []string{}
	}
	return stmts
}

var readOnlyRe = regexp.MustCompile(`(?i)^(SELECT|WITH|VALUES|EXPLAIN|PRAGMA)\b`)

// IsReadOnlyStatement reports whether stmt begins with a read-only keyword
// (SELECT/WITH/VALUES/EXPLAIN/PRAGMA). Such statements are safe to run against
// a snapshot without mutating state.
func IsReadOnlyStatement(stmt string) bool {
	return readOnlyRe.MatchString(strings.TrimSpace(stmt))
}

// FormatScalar renders one cell value as the textual form used in tables/CSV:
// nil -> "NULL", bool/number -> native string, string -> verbatim, anything
// else -> JSON.
func FormatScalar(v any) string {
	if v == nil {
		return "NULL"
	}
	switch val := v.(type) {
	case bool:
		if val {
			return "true"
		}
		return "false"
	case float64:
		// JSON-decoded numbers arrive as float64. Render without trailing
		// fractional noise when the value is a whole number, mirroring JS
		// Number -> String formatting.
		if val == float64(int64(val)) {
			return jsonInt(int64(val))
		}
		return jsonFloat(val)
	case int:
		return jsonInt(int64(val))
	case int64:
		return jsonInt(val)
	case float32:
		return jsonFloat(float64(val))
	case string:
		return val
	default:
		b, err := json.Marshal(val)
		if err != nil {
			return "NULL"
		}
		return string(b)
	}
}

// jsonInt reproduces JavaScript's String(number) for integers (no grouping,
// no exponent). strconv.FormatInt(base 10) gives exactly that.
func jsonInt(i int64) string { return strconv.FormatInt(i, 10) }

// jsonFloat reproduces JavaScript's shortest round-trip float formatting.
// 'g' with -1 precision picks the shortest form that round-trips.
func jsonFloat(f float64) string { return strconv.FormatFloat(f, 'g', -1, 64) }

var needsCsvQuoting = regexp.MustCompile(`[,"\n\r]`)

// csvField formats one cell per RFC-4180-ish CSV: quote if it contains a
// comma, double quote, or newline; double any embedded double quotes.
func csvField(v any) string {
	s := FormatScalar(v)
	if needsCsvQuoting.MatchString(s) {
		return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
	}
	return s
}

// RowsToCSV serializes a result set to CSV (header + one line per record,
// trailing LF).
func RowsToCSV(columns []string, rows [][]any) string {
	var lines []string
	colLine := make([]string, len(columns))
	for i, c := range columns {
		colLine[i] = csvField(c)
	}
	lines = append(lines, strings.Join(colLine, ","))

	for _, row := range rows {
		parts := make([]string, len(row))
		for i, cell := range row {
			parts[i] = csvField(cell)
		}
		lines = append(lines, strings.Join(parts, ","))
	}
	return strings.Join(lines, "\n") + "\n"
}

// RowsToMarkdown serializes a result set as a GitHub-flavored Markdown table.
// Pipes inside cells are escaped with a backslash so they don't break layout.
func RowsToMarkdown(columns []string, rows [][]any) string {
	esc := func(s string) string { return strings.ReplaceAll(s, "|", `\|`) }

	colCells := make([]string, len(columns))
	for i, c := range columns {
		colCells[i] = esc(c)
	}
	header := "| " + strings.Join(colCells, " | ") + " |"
	sepCells := make([]string, len(columns))
	for i := range columns {
		sepCells[i] = "---"
	}
	sep := "| " + strings.Join(sepCells, " | ") + " |"

	lines := []string{header, sep}
	for _, row := range rows {
		cells := make([]string, len(row))
		for i, cell := range row {
			cells[i] = esc(FormatScalar(cell))
		}
		lines = append(lines, "| "+strings.Join(cells, " | ")+" |")
	}
	return strings.Join(lines, "\n") + "\n"
}

// RowsToJSON serializes a result set as a pretty-printed JSON array of objects
// keyed by column name. Missing cells (row shorter than columns) become null.
func RowsToJSON(columns []string, rows [][]any) string {
	out := make([]map[string]any, 0, len(rows))
	for _, row := range rows {
		obj := make(map[string]any, len(columns))
		for i, col := range columns {
			if i < len(row) {
				obj[col] = row[i]
			} else {
				obj[col] = nil
			}
		}
		out = append(out, obj)
	}
	b, _ := json.MarshalIndent(out, "", "  ")
	return string(b)
}

var constraintRe = regexp.MustCompile(`(?i)^(PRIMARY\s+KEY|FOREIGN\s+KEY|UNIQUE|CHECK|CONSTRAINT)\b`)
var columnNameRe = regexp.MustCompile("^[`\"']?(\\w+)[`\"']?")

// extractColumnName pulls the column name (lower-cased) from a single column
// definition. Returns "" for table-level constraint lines (PRIMARY KEY,
// FOREIGN KEY, UNIQUE, CHECK, CONSTRAINT), which are not columns.
func extractColumnName(def string) string {
	if def == "" {
		return ""
	}
	if constraintRe.MatchString(def) {
		return ""
	}
	m := columnNameRe.FindStringSubmatch(def)
	if m == nil {
		return ""
	}
	return strings.ToLower(m[1])
}

var headerRe = regexp.MustCompile(`(?i)CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?[` + "`" + `\"']?(\w+)[` + "`" + `\"']?\s*\(`)

// SummarizeSchema best-effort extracts { table, columns } summaries from a
// CREATE TABLE script. Handles IF NOT EXISTS, quoted identifiers, and
// parenthesised type/constraint bodies (depth-tracked so a comma inside
// NUMERIC(10,2) doesn't split a column). Returns an empty slice on any error.
func SummarizeSchema(createSql string) []SchemaTable {
	// Defer a recover so a malformed input can never panic the caller.
	defer func() {
		_ = recover()
	}()

	results := []SchemaTable{}
	idxs := headerRe.FindAllStringSubmatchIndex(createSql, -1)
	if idxs == nil {
		return results
	}

	for _, loc := range idxs {
		// loc[0]=match start, loc[1]=match end, loc[2:4]=group1 (table name) span.
		table := strings.ToLower(createSql[loc[2]:loc[3]])
		startIdx := loc[1]

		// Walk forward to find the matching close paren of the table body.
		depth := 1
		i := startIdx
		for i < len(createSql) && depth > 0 {
			switch createSql[i] {
			case '(':
				depth++
			case ')':
				depth--
			}
			i++
		}
		if depth != 0 {
			// Unbalanced parens -> malformed; skip.
			continue
		}

		body := createSql[startIdx : i-1]
		var columns []string

		// Split body on top-level commas only.
		colDepth := 0
		current := strings.Builder{}
		for _, ch := range body {
			switch ch {
			case '(':
				colDepth++
				current.WriteRune(ch)
			case ')':
				colDepth--
				current.WriteRune(ch)
			case ',':
				if colDepth == 0 {
					if col := extractColumnName(strings.TrimSpace(current.String())); col != "" {
						columns = append(columns, col)
					}
					current.Reset()
				} else {
					current.WriteRune(ch)
				}
			default:
				current.WriteRune(ch)
			}
		}
		if col := extractColumnName(strings.TrimSpace(current.String())); col != "" {
			columns = append(columns, col)
		}

		if columns == nil {
			columns = []string{}
		}
		results = append(results, SchemaTable{Table: table, Columns: columns})
	}

	return results
}