llmgw/internal/store/dialect.go

package store

import (
	"context"
	"database/sql"
	"strconv"
	"strings"
)

// Dialect is which SQL database is underneath.
//
// There is one set of queries, not one per database, because the SQL this
// store writes is almost entirely portable — the only construct that genuinely
// differs is the placeholder, and `ON CONFLICT` upserts work in both. Two
// copies of every query would drift, and the drift would be discovered one
// query at a time in production.
type Dialect string

// The supported databases.
const (
	SQLiteDialect   Dialect = "sqlite"
	PostgresDialect Dialect = "postgres"
)

// rebind rewrites `?` placeholders for the dialect in use.
//
// SQLite takes `?`, PostgreSQL takes `$1`, `$2` and so on. Queries are written
// with `?` because that is what the majority of them were written with
// already, and rewriting is cheaper and safer than maintaining two spellings
// of each.
//
// A `?` inside a string literal would be rewritten too, which would be a bug —
// so this skips over quoted sections rather than scanning blindly. No query
// here contains one today; this exists so that adding one is not a silent
// corruption.
func (d Dialect) rebind(query string) string {
	if d != PostgresDialect || !strings.Contains(query, "?") {
		return query
	}
	var b strings.Builder
	b.Grow(len(query) + 8)

	n := 0
	inSingle, inDouble := false, false
	for i := 0; i < len(query); i++ {
		c := query[i]
		// A line comment is copied as it is. Read as SQL, the apostrophe in
		// "-- the key's owner" opened a string that swallowed every
		// placeholder after it, and PostgreSQL refused the query.
		if c == '-' && !inSingle && !inDouble && i+1 < len(query) && query[i+1] == '-' {
			end := strings.IndexByte(query[i:], '\n')
			if end < 0 {
				end = len(query) - i
			}
			b.WriteString(query[i : i+end])
			i += end - 1
			continue
		}
		switch {
		case c == '\'' && !inDouble:
			inSingle = !inSingle
		case c == '"' && !inSingle:
			inDouble = !inDouble
		case c == '?' && !inSingle && !inDouble:
			n++
			b.WriteByte('$')
			b.WriteString(strconv.Itoa(n))
			continue
		}
		b.WriteByte(c)
	}
	return b.String()
}

// conn is the part of *sql.DB and *sql.Tx this store uses.
//
// Named so the same query helpers work inside and outside a transaction, and
// so every path goes through the rebinding wrapper rather than some of them
// reaching the driver directly with the wrong placeholders.
type conn interface {
	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
	QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
	// PrepareContext is here because a prepared statement carries the query
	// with it: rebinding at Exec time would be too late, the placeholders are
	// fixed when it is prepared.
	PrepareContext(ctx context.Context, query string) (*sql.Stmt, error)
}

// bound is a conn that rewrites placeholders on the way through.
//
// A wrapper rather than a change at each call site: there are sixty of those,
// and sixty chances to forget one. Here the only way to reach the driver is
// through the rewrite.
type bound struct {
	c Dialect
	// to is the underlying *sql.DB or *sql.Tx.
	to conn
}

func (b bound) ExecContext(ctx context.Context, q string, args ...any) (sql.Result, error) {
	return b.to.ExecContext(ctx, b.c.rebind(q), args...)
}

func (b bound) QueryContext(ctx context.Context, q string, args ...any) (*sql.Rows, error) {
	return b.to.QueryContext(ctx, b.c.rebind(q), args...)
}

func (b bound) QueryRowContext(ctx context.Context, q string, args ...any) *sql.Row {
	return b.to.QueryRowContext(ctx, b.c.rebind(q), args...)
}

func (b bound) PrepareContext(ctx context.Context, q string) (*sql.Stmt, error) {
	return b.to.PrepareContext(ctx, b.c.rebind(q))
}