Test Text-to-SQL Queries Against a Real Postgres

sbx new --name sql-seed --no-env
id=sb_v9na1yk0k4    # the id that sbx new printed

sbx exec $id --timeout 300 "docker run -d --name pg --restart unless-stopped \
  -e POSTGRES_PASSWORD=pw -p 127.0.0.1:5432:5432 postgres:16"
sbx exec $id "docker exec pg pg_isready -U postgres"

Status: Downloaded newer image for postgres:16
/var/run/postgresql:5432 - accepting connections

CREATE DATABASE shop;
\c shop

CREATE TABLE customers (id int PRIMARY KEY, name text NOT NULL, country text NOT NULL, created_at date NOT NULL);
CREATE TABLE products (id int PRIMARY KEY, name text NOT NULL, category text NOT NULL, price_cents int NOT NULL);
CREATE TABLE orders (
  id int PRIMARY KEY, customer_id int NOT NULL REFERENCES customers(id), ordered_at date NOT NULL,
  status text NOT NULL CHECK (status IN ('paid', 'refunded'))
);
CREATE TABLE order_items (
  order_id int NOT NULL REFERENCES orders(id), product_id int NOT NULL REFERENCES products(id),
  quantity int NOT NULL, PRIMARY KEY (order_id, product_id)
);

INSERT INTO customers VALUES (1, 'Ada', 'US', '2026-01-04'), (2, 'Bruno', 'DE', '2026-02-11'),
  (3, 'Chen', 'US', '2026-03-02'), (4, 'Dalia', 'FR', '2026-05-19'), (5, 'Emeka', 'DE', '2026-07-23');
INSERT INTO products VALUES (1, 'Keyboard', 'hardware', 8900), (2, 'Mouse', 'hardware', 2900),
  (3, 'Monitor', 'hardware', 24900), (4, 'IDE seat', 'software', 12000), (5, 'VPN year', 'software', 6000);
INSERT INTO orders VALUES (101, 1, '2026-08-02', 'paid'), (102, 2, '2026-08-15', 'paid'),
  (103, 1, '2026-09-01', 'refunded'), (104, 3, '2026-09-09', 'paid'), (105, 4, '2026-09-21', 'paid'),
  (106, 5, '2026-10-03', 'paid'), (107, 2, '2026-10-12', 'paid');
INSERT INTO order_items VALUES (101, 1, 1), (101, 2, 2), (102, 3, 2), (103, 4, 1), (104, 4, 3),
  (104, 5, 1), (105, 2, 1), (106, 3, 1), (106, 1, 1), (107, 5, 4);

-- The role the agent's queries run as: it can read, nothing else.
CREATE ROLE agent_ro LOGIN PASSWORD 'ro';
REVOKE ALL ON SCHEMA public FROM PUBLIC;
GRANT CONNECT ON DATABASE shop TO agent_ro;
GRANT USAGE ON SCHEMA public TO agent_ro;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO agent_ro;
ALTER ROLE agent_ro SET default_transaction_read_only = on;
ALTER ROLE agent_ro SET statement_timeout = '5s';

sbx scp ./seed.sql $id:/home/user/
sbx exec $id "docker exec -i pg psql -U postgres -v ON_ERROR_STOP=1 < seed.sql"

CREATE DATABASE
You are now connected to database "shop" as user "postgres".
CREATE TABLE
…
INSERT 0 10
CREATE ROLE
…
ALTER ROLE

import json

PSQL = "docker exec -i pg psql -X -q --csv -v ON_ERROR_STOP=1 -U agent_ro -d shop"
MAX_ROWS = 50


def make_run_sql(sandbox):
    """Return a run_sql tool bound to one sandbox that has the seeded Postgres."""

    def run_sql(query: str) -> str:
        """Run one SQL query against the shop database, read-only.
        Returns {"ok": true, "columns": [...], "rows": [CSV lines]} or {"ok": false, "error": "..."}."""
        sandbox.files.write("query.sql", f"BEGIN TRANSACTION READ ONLY;\n{query}\n;\nROLLBACK;\n")
        r = sandbox.exec(f"{PSQL} < query.sql", timeout_seconds=30)
        if r.timed_out:
            return json.dumps({"ok": False, "error": "query took longer than 30 seconds"})
        if r.exit_code != 0:
            return json.dumps({"ok": False, "error": r.stderr.strip()})
        lines = r.stdout.strip().splitlines()
        result = {"ok": True, "columns": lines[0].split(",") if lines else [], "rows": lines[1 : MAX_ROWS + 1]}
        if len(lines) - 1 > MAX_ROWS:
            result["truncated"] = True
        return json.dumps(result)

    return run_sql

run_sql("""SELECT table_name, string_agg(column_name || ' ' || data_type, ', ' ORDER BY ordinal_position)
           FROM information_schema.columns WHERE table_schema = 'public'
           GROUP BY table_name ORDER BY table_name""")

{"ok": true, "columns": ["table_name", "columns"], "rows": ["customers,\"id integer, name text, country text, created_at date\"", "order_items,\"order_id integer, product_id integer, quantity integer\"", "orders,\"id integer, customer_id integer, ordered_at date, status text\"", "products,\"id integer, name text, category text, price_cents integer\""]}