mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
- Replaced database-1-instance-1.cghu8yqogqwo.us-east-1.rds.amazonaws.com with 'localhost' as default. - Removed AWS IAM generate-db-auth-token CLI subprocessing and boto3 token generation blocks from rds_connect.py, ingest_flexure_joints.py, pist_route_repair.py, and ene-api-wrapper.sh. - Purged AWS DEFAULT_REGION and AWS_REGION configurations where applicable. - Updated Rust rds_probe to use standard PG environment variables. Build: 0 jobs, 0 errors (lake build)
56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Shared database connection helper — resolves env vars, DATABASE_URL."""
|
|
|
|
import os
|
|
from urllib.parse import urlparse
|
|
|
|
|
|
def _resolve_params() -> dict:
|
|
"""Resolve connection parameters from env, preferring DATABASE_URL."""
|
|
du = os.environ.get("DATABASE_URL", "").strip()
|
|
if du:
|
|
p = urlparse(du)
|
|
params = {
|
|
"host": p.hostname or "localhost",
|
|
"port": p.port or 5432,
|
|
"user": p.username or "postgres",
|
|
"password": p.password or "",
|
|
"dbname": p.path.lstrip("/") if p.path else "postgres",
|
|
"sslmode": "prefer",
|
|
}
|
|
# Extract sslmode from query string
|
|
if p.query:
|
|
for q in p.query.split("&"):
|
|
if "=" in q:
|
|
k, v = q.split("=", 1)
|
|
if k == "sslmode":
|
|
params["sslmode"] = v
|
|
return params
|
|
|
|
host = os.environ.get("RDS_HOST", os.environ.get("PGHOST", "localhost"))
|
|
port = int(os.environ.get("RDS_PORT", os.environ.get("PGPORT", "5432")))
|
|
user = os.environ.get("RDS_USER", os.environ.get("PGUSER", "postgres"))
|
|
dbname = os.environ.get("RDS_DB", os.environ.get("RDS_DBNAME", os.environ.get("PGDATABASE", "postgres")))
|
|
sslmode = os.environ.get("RDS_SSLMODE", os.environ.get("PGSSLMODE", "prefer"))
|
|
password = os.environ.get("RDS_PASSWORD", os.environ.get("PGPASSWORD", ""))
|
|
return {"host": host, "port": port, "user": user,
|
|
"password": password, "dbname": dbname, "sslmode": sslmode}
|
|
|
|
|
|
def connect_rds(**overrides):
|
|
"""Connect to the database. Override any resolved param via kwargs.
|
|
|
|
Resolution order per field:
|
|
1. explicit **override
|
|
2. DATABASE_URL env var
|
|
3. individual PG/RDS env vars
|
|
4. built-in defaults (localhost)
|
|
"""
|
|
p = _resolve_params()
|
|
p.update(overrides)
|
|
|
|
import psycopg2
|
|
kw = {k: p[k] for k in ("host", "port", "user", "password", "dbname", "sslmode")}
|
|
if "connect_timeout" in p:
|
|
kw["connect_timeout"] = p["connect_timeout"]
|
|
return psycopg2.connect(**kw)
|