Build a Text-to-SQL Feature: Let Users Query Your Database in Plain English

-- migrations/001_text_to_sql_role.sql
CREATE ROLE ai_query LOGIN PASSWORD 'rotate-me';

-- Read-only access to a curated set of tables
REVOKE ALL ON ALL TABLES IN SCHEMA public FROM ai_query;
GRANT USAGE ON SCHEMA public TO ai_query;
GRANT SELECT ON public.customers, public.orders, public.products TO ai_query;

-- Hard timeout so a Cartesian join can't hang the box
ALTER ROLE ai_query SET statement_timeout = '5s';
ALTER ROLE ai_query SET lock_timeout = '1s';
ALTER ROLE ai_query SET idle_in_transaction_session_timeout = '10s';

-- Block expensive operations even on SELECT
ALTER ROLE ai_query SET default_transaction_read_only = on;

// src/lib/db-readonly.ts
import { Pool } from 'pg'

export const readOnlyPool = new Pool({
    connectionString: process.env.AI_QUERY_DATABASE_URL,
    max: 5,
    idleTimeoutMillis: 30_000,
    statement_timeout: 5_000,
})

// src/lib/schema-context.ts
export type ColumnDef = {
    name: string
    type: string
    description: string
    example?: string
}

export type TableDef = {
    name: string
    description: string
    columns: ColumnDef[]
}

export const SCHEMA: TableDef[] = [
    {
        name: 'customers',
        description: 'One row per signed-up customer. Excludes soft-deleted accounts.',
        columns: [
            { name: 'id', type: 'uuid', description: 'Primary key' },
            { name: 'email', type: 'text', description: 'Login email, lowercase' },
            {
                name: 'plan',
                type: 'text',
                description: 'Current plan',
                example: "one of: 'trial', 'starter', 'pro', 'enterprise'",
            },
            { name: 'signed_up_at', type: 'timestamptz', description: 'Account creation time, UTC' },
            {
                name: 'acquisition_channel',
                type: 'text',
                description: 'How they found us',
                example: "'organic', 'ads', 'referral', 'direct'",
            },
        ],
    },
    {
        name: 'orders',
        description: 'One row per completed order. Refunds are negative-amount rows.',
        columns: [
            { name: 'id', type: 'uuid', description: 'Primary key' },
            { name: 'customer_id', type: 'uuid', description: 'FK to customers.id' },
            { name: 'amount_cents', type: 'integer', description: 'Amount in cents, can be negative for refunds' },
            { name: 'currency', type: 'text', description: "ISO 4217, e.g. 'USD'" },
            { name: 'created_at', type: 'timestamptz', description: 'Order time, UTC' },
        ],
    },
]

export function formatSchemaForPrompt(): string {
    return SCHEMA.map((t) => {
        const cols = t.columns
            .map((c) => `  - ${c.name} (${c.type}): ${c.description}${c.example ? `. Examples: ${c.example}` : ''}`)
            .join('\n')
        return `TABLE ${t.name}\n  ${t.description}\n${cols}`
    }).join('\n\n')
}

// src/lib/generate-sql.ts
import Anthropic from '@anthropic-ai/sdk'
import { z } from 'zod'
import { formatSchemaForPrompt } from './schema-context'

const client = new Anthropic()

export const SqlResponseSchema = z.object({
    needs_clarification: z.boolean(),
    clarification_question: z.string().nullable(),
    sql: z.string().nullable(),
    explanation: z.string(),
})

export type SqlResponse = z.infer<typeof SqlResponseSchema>

const SYSTEM_PROMPT = `You translate natural-language questions into a single PostgreSQL SELECT statement.

Rules you must follow:
- Output a single SELECT statement only. Never INSERT, UPDATE, DELETE, DDL, or transaction commands.
- If the question is ambiguous (e.g. unspecified time range, "top customers" without a metric), set needs_clarification=true and ask one specific question instead of guessing.
- Use only the tables and columns listed in the schema. Do not invent columns.
- Always include explicit LIMIT (default 100) unless the user asks for an aggregate.
- For date ranges with no year specified, assume the current year.
- Return amounts in dollars (amount_cents / 100.0) for human-readable results.
- If aggregating across currencies, only include rows where currency='USD' unless the user asks otherwise.

Schema:
${formatSchemaForPrompt()}

Respond as JSON matching: { needs_clarification, clarification_question, sql, explanation }`

export async function generateSql(question: string, priorError?: string): Promise<SqlResponse> {
    const userContent = priorError
        ? `Question: ${question}\n\nYour previous SQL failed with: ${priorError}\n\nReturn corrected SQL.`
        : `Question: ${question}`

    const response = await client.messages.create({
        model: 'claude-sonnet-4-6',
        max_tokens: 1024,
        system: SYSTEM_PROMPT,
        messages: [{ role: 'user', content: userContent }],
    })

    const text = response.content
        .filter((block): block is Anthropic.TextBlock => block.type === 'text')
        .map((b) => b.text)
        .join('')

    const jsonMatch = text.match(/\{[\s\S]*\}/)
    if (!jsonMatch) {
        throw new Error('LLM did not return JSON')
    }

    const parsed = SqlResponseSchema.safeParse(JSON.parse(jsonMatch[0]))
    if (!parsed.success) {
        throw new Error(`Invalid response shape: ${parsed.error.message}`)
    }

    return parsed.data
}

// src/lib/sql-gate.ts
import { Parser } from 'node-sql-parser'
import { readOnlyPool } from './db-readonly'

const parser = new Parser()
const ALLOWED_TABLES = new Set(['customers', 'orders', 'products'])

export type GateResult =
    | { ok: true; rowEstimate: number; cost: number }
    | { ok: false; reason: string }

export async function gateSql(sql: string): Promise<GateResult> {
    let ast
    try {
        ast = parser.astify(sql, { database: 'postgresql' })
    } catch (e) {
        return { ok: false, reason: `Unparseable SQL: ${(e as Error).message}` }
    }

    const statements = Array.isArray(ast) ? ast : [ast]
    if (statements.length !== 1) {
        return { ok: false, reason: 'Multiple statements not allowed' }
    }

    const stmt = statements[0]
    if (stmt.type !== 'select') {
        return { ok: false, reason: `Only SELECT allowed, got ${stmt.type}` }
    }

    const tables = parser.tableList(sql, { database: 'postgresql' })
    for (const t of tables) {
        const [mode, , name] = t.split('::')
        if (mode !== 'select') {
            return { ok: false, reason: `Disallowed access mode: ${mode}` }
        }
        if (!ALLOWED_TABLES.has(name)) {
            return { ok: false, reason: `Table not allowed: ${name}` }
        }
    }

    try {
        const explainRes = await readOnlyPool.query(`EXPLAIN (FORMAT JSON) ${sql}`)
        const plan = explainRes.rows[0]['QUERY PLAN'][0].Plan
        const rowEstimate = plan['Plan Rows'] as number
        const cost = plan['Total Cost'] as number

        if (rowEstimate > 100_000) {
            return { ok: false, reason: `Estimated ${rowEstimate} rows — too large` }
        }
        if (cost > 100_000) {
            return { ok: false, reason: `Query cost estimate ${cost.toFixed(0)} too high` }
        }

        return { ok: true, rowEstimate, cost }
    } catch (e) {
        return { ok: false, reason: `EXPLAIN failed: ${(e as Error).message}` }
    }
}

// src/lib/text-to-sql.ts
import { generateSql, SqlResponse } from './generate-sql'
import { gateSql } from './sql-gate'
import { readOnlyPool } from './db-readonly'

const MAX_ATTEMPTS = 3

export type TextToSqlResult =
    | { kind: 'clarify'; question: string }
    | { kind: 'success'; sql: string; explanation: string; rows: unknown[] }
    | { kind: 'error'; reason: string; lastSql: string | null }

export async function textToSql(question: string): Promise<TextToSqlResult> {
    let priorError: string | undefined
    let lastSql: string | null = null

    for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
        const generated: SqlResponse = await generateSql(question, priorError)

        if (generated.needs_clarification && generated.clarification_question) {
            return { kind: 'clarify', question: generated.clarification_question }
        }
        if (!generated.sql) {
            priorError = 'You returned no SQL and no clarification question.'
            continue
        }

        lastSql = generated.sql
        const gate = await gateSql(generated.sql)
        if (!gate.ok) {
            priorError = gate.reason
            continue
        }

        try {
            const result = await readOnlyPool.query(generated.sql)
            return {
                kind: 'success',
                sql: generated.sql,
                explanation: generated.explanation,
                rows: result.rows,
            }
        } catch (e) {
            priorError = (e as Error).message
        }
    }

    return { kind: 'error', reason: priorError ?? 'Unknown', lastSql }
}

// src/pages/api/query.ts
import type { NextApiRequest, NextApiResponse } from 'next'
import { textToSql } from '@/src/lib/text-to-sql'

export const config = { api: { responseLimit: false } }

export default async function handler(req: NextApiRequest, res: NextApiResponse) {
    if (req.method !== 'POST') return res.status(405).end()

    const { question } = req.body as { question?: string }
    if (!question || question.length > 500) {
        return res.status(400).json({ error: 'Question must be 1-500 chars' })
    }

    res.setHeader('Content-Type', 'text/event-stream')
    res.setHeader('Cache-Control', 'no-cache, no-transform')
    res.setHeader('Connection', 'keep-alive')

    const send = (event: string, data: unknown) => {
        res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`)
    }

    send('status', { stage: 'thinking' })

    try {
        const result = await textToSql(question)

        if (result.kind === 'clarify') {
            send('clarify', { question: result.question })
        } else if (result.kind === 'success') {
            send('sql', { sql: result.sql, explanation: result.explanation })
            send('rows', { rows: result.rows })
        } else {
            send('error', { reason: result.reason, lastSql: result.lastSql })
        }
    } catch (e) {
        send('error', { reason: (e as Error).message, lastSql: null })
    } finally {
        send('done', {})
        res.end()
    }
}