mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
Creates 4-Infrastructure/shim/rds_connect.py with a single connect_rds()
function that resolves connection parameters in priority order:
1. explicit kwargs
2. DATABASE_URL env var (postgres://user:pass@host:port/dbname?sslmode=...)
3. individual RDS_* env vars (RDS_HOST, RDS_PORT, RDS_USER, etc.)
4. built-in defaults
Auth resolution (when password is empty or RDS_IAM=1):
1. RDS_IAM_TOKEN env var (pre-computed)
2. boto3 SDK generate_db_auth_token (preferred)
3. subprocess aws rds generate-db-auth-token (fallback)
4. RDS_PASSWORD env var (non-IAM)
Replaces 8 connection pattern variants across 14 active shims:
- subprocess + RDS_IAM_TOKEN fallback: pist_trace_classify_mcp, joint_classifier,
pist_prove_and_classify, ingest_57_flexures
- boto3 SDK: ene_wiki_body_reingest, ene_migrate_and_tag, dataset_ingest_rds
- subprocess + RDS_PASSWORD: batch_embed_artifacts, sync_wiki_to_rds, seed_flexure_dataset
- RDS_IAM_AUTH: pist_classify
- bashrc parsed: credential_loader
v1.4a benchmark confirmed at 100% after refactor.
65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os, sys, json
|
|
from rds_connect import connect_rds
|
|
|
|
RDS_HOST = None
|
|
RDS_USER = None
|
|
|
|
def _init_rds():
|
|
global RDS_HOST, RDS_USER
|
|
bashrc = os.path.expanduser('~/.bashrc')
|
|
if os.path.exists(bashrc):
|
|
with open(bashrc) as f:
|
|
for line in f:
|
|
if line.startswith('export RDS_HOST='):
|
|
RDS_HOST = line.split('"')[1] if '"' in line else line.split('=')[1].strip()
|
|
elif line.startswith('export RDS_USER='):
|
|
RDS_USER = line.split('"')[1] if '"' in line else line.split('=')[1].strip()
|
|
if not RDS_HOST or not RDS_USER:
|
|
raise RuntimeError("RDS_HOST and RDS_USER must be set in ~/.bashrc")
|
|
|
|
def load_credential(pkg: str, password: str = None) -> str:
|
|
_init_rds()
|
|
conn = connect_rds(host=RDS_HOST, user=RDS_USER, dbname='postgres')
|
|
cur = conn.cursor()
|
|
if password is None:
|
|
password = RDS_HOST
|
|
cur.execute(
|
|
"SELECT pgp_sym_decrypt(encrypted_payload, %s) FROM credential_store.credentials WHERE pkg = %s",
|
|
(password, pkg))
|
|
row = cur.fetchone()
|
|
cur.close()
|
|
conn.close()
|
|
if row is None:
|
|
raise KeyError(f"Credential '{pkg}' not found")
|
|
return row[0]
|
|
|
|
def list_credentials() -> list[dict]:
|
|
_init_rds()
|
|
conn = connect_rds(host=RDS_HOST, user=RDS_USER, dbname='postgres')
|
|
cur = conn.cursor()
|
|
cur.execute(
|
|
"SELECT id, pkg, provider, classification, created_at, is_active "
|
|
"FROM credential_store.credentials WHERE is_active = true ORDER BY pkg")
|
|
rows = cur.fetchall()
|
|
cur.close()
|
|
conn.close()
|
|
return [
|
|
{'id': r[0], 'pkg': r[1], 'provider': r[2], 'classification': r[3],
|
|
'created_at': str(r[4]), 'active': r[5]}
|
|
for r in rows]
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) < 2:
|
|
creds = list_credentials()
|
|
print("Available credentials:")
|
|
for c in creds:
|
|
print(f" {c['pkg']:30s} provider={c['provider']:15s} active={c['active']}")
|
|
sys.exit(0)
|
|
if sys.argv[1] == '--export' and len(sys.argv) >= 3:
|
|
key = load_credential(sys.argv[2])
|
|
print(f"export {sys.argv[2].upper().replace('-', '_')}='{key}'")
|
|
else:
|
|
key = load_credential(sys.argv[1])
|
|
print(key)
|