postgres-connect-and-ping

copy

import sys
import os
import subprocess
import time
import urllib.request
import json

# ─────────────────────────────────────────
# PRE_EXECUTION
# ─────────────────────────────────────────

# 1. Fetch registry entry
REGISTRY_URL = "https://checklist.day/api/registry/psycopg2-binary"

for attempt in range(2):
    try:
        req = urllib.request.Request(REGISTRY_URL, headers={"User-Agent": "checklist-agent/1.0"})
        with urllib.request.urlopen(req, timeout=10) as resp:
            registry = json.loads(resp.read())
            break
    except Exception as e:
        if attempt == 1:
            print(f"ABORT: registry unreachable after 2 attempts — {e}")
            sys.exit(1)
        time.sleep(2)

# 2. Surface warnings
warnings = registry.get("warnings", [])
if warnings:
    print("[psycopg2-binary] WARNINGS:")
    for w in warnings if isinstance(warnings, list) else [warnings]:
        print(f"  ⚠ {w}")

# ─────────────────────────────────────────
# EXECUTION
# ─────────────────────────────────────────

# 3. Install
# FOOTGUN: use psycopg2-binary not psycopg2 — binary includes libpq, no build tools needed
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "psycopg2-binary>=2.9.9"])

import psycopg2
from psycopg2 import OperationalError

# 4. Validate inputs
DATABASE_URL = os.environ.get("DATABASE_URL")
if not DATABASE_URL:
    print("ABORT: DATABASE_URL env var not set")
    print("  format: postgresql://user:pass@host:5432/dbname")
    sys.exit(1)

# FOOTGUN: psycopg2 does not accept postgres:// scheme — must be postgresql://
if DATABASE_URL.startswith("postgres://"):
    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)
    print("  [fixed] postgres:// → postgresql://")

# 5. Connect and ping
connected = False
server_version = None
latency_ms = None
error = None
conn = None

try:
    t0 = time.perf_counter()
    conn = psycopg2.connect(DATABASE_URL, connect_timeout=10)
    latency_ms = round((time.perf_counter() - t0) * 1000, 1)

    # FOOTGUN: use a cursor to actually verify the connection is usable, not just opened
    with conn.cursor() as cur:
        cur.execute("SELECT version()")
        server_version = cur.fetchone()[0]

    connected = True
    print(f"  connected in {latency_ms}ms")
    print(f"  server: {server_version}")

except OperationalError as e:
    # Surface exact error — agents need the real message to diagnose
    error = str(e).strip()
    print(f"  connection failed: {error}")

finally:
    if conn:
        conn.close()

# ─────────────────────────────────────────
# POST_EXECUTION
# ─────────────────────────────────────────

result = {
    "connected":      connected,
    "server_version": server_version,
    "latency_ms":     latency_ms,
    "error":          error,
}
print(json.dumps(result, indent=2))

if not connected:
    print(f"FAIL: could not connect — {error}")
    sys.exit(1)

assert server_version is not None, "FAIL: connected but server_version is null"
assert latency_ms is not None, "FAIL: connected but latency_ms is null"
assert latency_ms < 5000, f"FAIL: latency {latency_ms}ms exceeds 5s threshold"

print("PASS")