Skip to content
Collimate

Cookbook: live Postgres

03 / Live service

For DB-agent and text-to-SQL RL, the setup cost is standing up a database: initdb, starting the server, creating the schema, and loading fixtures — paid before a single query runs. Bake a live-service ready state: initialize and start Postgres at bake time and leave it running, so every fork wakes with the database already up and a rollout just runs queries against it.

The idea

Section titled “The idea”
bake:  boot guest → init + start Postgres → block until pg_isready → SNAPSHOT → publish
fork:  restore snapshot → Postgres is already serving → rollout runs queries

The warm command starts the server detached, then blocks until pg_isready reports the server is accepting connections — not merely that the process launched. See the warm-command contract.

The warm script

Section titled “The warm script”

/usr/local/bin/collimate-warm initializes a data directory (idempotently), starts the server detached, waits for it to accept connections, then loads the schema and fixtures your rollouts query:

#!/bin/sh
set -e


ip link set lo up 2>/dev/null || true
export PGDATA=/var/lib/postgresql/data


# 1. Initialize the cluster once. Skip if a data dir already exists so the script
#    is safe to re-run as a build-time self-test.
if [ ! -s "$PGDATA/PG_VERSION" ]; then
  su postgres -c "initdb -D $PGDATA --auth=trust" >/var/log/pg-init.log 2>&1
fi


# 2. Start the server DETACHED so it outlives this script and lands in the
#    snapshot. pg_ctl already backgrounds the postmaster; bind to localhost.
su postgres -c "pg_ctl -D $PGDATA -l /var/log/pg.log -o '-c listen_addresses=127.0.0.1' start"


# 3. Block until Postgres actually ACCEPTS CONNECTIONS — the real serving signal.
for i in $(seq 1 60); do
  if pg_isready -h 127.0.0.1 -q; then break; fi
  sleep 1
  [ "$i" = 60 ] && { echo pg-timeout >&2; exit 1; }
done


# 4. Initialize schema + fixtures once, so every fork wakes with the data loaded.
#    Guard with a sentinel so re-runs don't double-apply.
if ! psql -h 127.0.0.1 -U postgres -tAc "SELECT to_regclass('public.orders')" | grep -q orders; then
  psql -h 127.0.0.1 -U postgres -f /opt/schema.sql
  psql -h 127.0.0.1 -U postgres -f /opt/fixtures.sql
fi


echo pg-ready; exit 0

Bake it:

from collimate_rl import Ready, Template, connect


client = connect()   # a col_pro_* key via COLLIMATE_API_KEY


handle = Template.bake(
    client,
    name="postgres",
    image="ghcr.io/you/postgres@sha256:…",             # pin by digest
    ready=Ready(cmd="/usr/local/bin/collimate-warm"),  # launch detached → block → exit 0
    width=64,
)
handle.wait(timeout=600)

Run the warm script at build time as a self-test so a schema typo or a server that never comes up fails the build, not the bake.

Using it from the SDK

Section titled “Using it from the SDK”

A rollout attaches to the live, pre-seeded database and runs queries — the initdb / start / load cost is already paid:

from collimate_rl import connect


client = connect(api_key="col_...")


# The rollout body runs inside the sandbox. It connects to the Postgres the
# ready-state snapshot is already running — no per-fork database startup.
ROLLOUT = r"""
import psycopg2


conn = psycopg2.connect(host="127.0.0.1", user="postgres", dbname="postgres")
cur = conn.cursor()


# A text-to-SQL agent would generate this query; grade it against expected rows.
cur.execute("SELECT customer_id, SUM(total) FROM orders GROUP BY customer_id ORDER BY 2 DESC LIMIT 3")
for row in cur.fetchall():
    print(row)


conn.close()
"""


sb = client.create_sandbox("postgres")
result = client.exec(sb["id"], commands=[["python3", "-c", ROLLOUT]], timeout_seconds=60)
print(result["stdout"])
client.delete_sandbox(sb["id"])

Fan a warmed parent out across many rollouts with fork — every child inherits the same live, seeded database, and each child's writes are private:

# count is the per-call fan-out, capped by tier (Pro 32, Ent 64); call again for more.
children = client.fork(sb["id"], count=32)["children"]
for child in children:
    client.exec(child["id"], commands=[["python3", "-c", ROLLOUT]], timeout_seconds=60)

Reach it from outside

Section titled “Reach it from outside”

The warm script above starts Postgres with -c listen_addresses=127.0.0.1, so it is reachable in-guest only — a connection arriving on the fork's NIC reaches nothing. To drive the database from a client on your machine — the connect-from-outside access model — start Postgres on all interfaces and expose the port:

  • In the warm script, set listen_addresses = '*' (e.g. pg_ctl … -o "-c listen_addresses=*") so the server binds the per-fork NIC too.
  • Bake it with ready=Ready(cmd="/usr/local/bin/collimate-warm", ports=[5432]) so the fork gets a guest NIC with 5432 exposed.

Then forward the port and point any Postgres client at fwd.host / fwd.port. fwd.url (an http://… address) does not apply to a wire protocol like Postgres — use the host and port:

import psycopg2
from collimate_rl import connect


client = connect(api_key="col_...")


with client.session("postgres") as sb:
    with sb.port_forward(5432) as fwd:
        # non-HTTP service: use fwd.host / fwd.port, not fwd.url
        conn = psycopg2.connect(host=fwd.host, port=fwd.port, user="postgres", dbname="postgres")
        # or a DSN: postgresql://postgres@{fwd.host}:{fwd.port}/postgres
        cur = conn.cursor()
        cur.execute("SELECT customer_id, SUM(total) FROM orders GROUP BY customer_id LIMIT 3")
        print(cur.fetchall())
        conn.close()
# leaving the `with` deletes the sandbox

Next steps

Section titled “Next steps”