Jonathan
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
import os
|
|
import time
|
|
|
|
import psycopg
|
|
|
|
|
|
def connect() -> psycopg.Connection:
|
|
dsn = os.environ["DATABASE_URL"]
|
|
return psycopg.connect(dsn)
|
|
|
|
|
|
def wait_for_db(timeout: float = 60.0, interval: float = 2.0) -> psycopg.Connection:
|
|
"""Block until the database is up and the `Job` table exists.
|
|
|
|
Repeatedly opens a connection and probes for the `Job` table's presence.
|
|
Tolerates the DB not being reachable yet (`psycopg.OperationalError`) and
|
|
migrations not having run yet (`psycopg.errors.UndefinedTable`), retrying
|
|
with a fixed `interval` until `timeout` elapses. Returns the open, usable
|
|
connection on success; re-raises the last error on timeout.
|
|
"""
|
|
deadline = time.monotonic() + timeout
|
|
printed_waiting = False
|
|
last_error: Exception | None = None
|
|
while time.monotonic() < deadline:
|
|
conn: psycopg.Connection | None = None
|
|
try:
|
|
conn = connect()
|
|
with conn.cursor() as cur:
|
|
cur.execute('SELECT 1 FROM "Job" LIMIT 0')
|
|
return conn
|
|
except (psycopg.OperationalError, psycopg.errors.UndefinedTable) as exc:
|
|
last_error = exc
|
|
if conn is not None:
|
|
conn.close()
|
|
if not printed_waiting:
|
|
print("worker: waiting for database schema...", flush=True)
|
|
printed_waiting = True
|
|
time.sleep(interval)
|
|
assert last_error is not None
|
|
raise last_error
|