gendesign/data/sql/32_load_domrf_raw.py
2026-04-26 22:30:52 +03:00

99 lines
3.9 KiB
Python

"""Universal loader: dumps every scraped JSON from data/raw/domrf_full/ into Postgres
as JSONB rows in domrf_raw_endpoints. Idempotent (UPSERT by snapshot+section+endpoint).
Usage:
PG_HOST=localhost PG_PORT=15432 PGPASSWORD=... python data/sql/32_load_domrf_raw.py
"""
import os, json, subprocess, sys, datetime
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, '..', 'raw', 'domrf_full')
PG_PORT = os.environ.get('PG_PORT', '5432')
PG_USER = os.environ.get('PG_USER', 'gendesign')
PG_DB = os.environ.get('PG_DB', 'gendesign')
PG_PASS = os.environ.get('PGPASSWORD', '')
SNAP = os.environ.get('SNAPSHOT_DATE', datetime.date.today().isoformat())
def psql(sql):
cmd = ['docker', 'run', '--rm', '-i', '-e', f'PGPASSWORD={PG_PASS}',
'postgres:16-alpine', 'psql', '-h', 'host.docker.internal', '-p', PG_PORT,
'-U', PG_USER, '-d', PG_DB, '-v', 'ON_ERROR_STOP=1', '--quiet']
res = subprocess.run(cmd, input=sql, capture_output=True, text=True, encoding='utf-8')
if res.returncode != 0:
print('psql FAILED:', res.stderr[-1500:], file=sys.stderr)
raise SystemExit(res.returncode)
return res.stdout
def psql_stdin(sql, stdin_data):
"""Use psql with COPY via stdin to load big JSONs (avoids cmdline length limits)."""
cmd = ['docker', 'run', '--rm', '-i', '-e', f'PGPASSWORD={PG_PASS}',
'postgres:16-alpine', 'psql', '-h', 'host.docker.internal', '-p', PG_PORT,
'-U', PG_USER, '-d', PG_DB, '-v', 'ON_ERROR_STOP=1', '--quiet']
res = subprocess.run(cmd, input=stdin_data, capture_output=True, text=True, encoding='utf-8')
if res.returncode != 0:
print('psql FAILED:', res.stderr[-1500:], file=sys.stderr)
raise SystemExit(res.returncode)
return res.stdout
def main():
schema = open(os.path.join(HERE, '31_schema_domrf_raw.sql'), encoding='utf-8').read()
psql(schema)
print('schema applied')
loaded = 0
skipped = 0
by_section = {}
for section in sorted(os.listdir(ROOT)):
sec_dir = os.path.join(ROOT, section)
if not os.path.isdir(sec_dir):
continue
files = [f for f in os.listdir(sec_dir) if f.endswith('.json')]
for fname in sorted(files):
fpath = os.path.join(sec_dir, fname)
try:
text = open(fpath, encoding='utf-8').read()
json.loads(text) # validate
except Exception as e:
print(f' SKIP invalid JSON: {section}/{fname}: {e}')
skipped += 1
continue
endpoint = fname[:-5] # strip .json
# Use COPY ... FROM STDIN for safe binary insertion
# Use psycopg-style escape: \\ → \\\\, replace tabs/newlines in JSON keys
# Easier: pass JSON via stdin to a Python heredoc-style INSERT
sql = f"""
INSERT INTO domrf_raw_endpoints (snapshot_date, section, endpoint, source_url, payload, payload_size, fetched_at)
VALUES ('{SNAP}', $sec${section}$sec$, $end${endpoint}$end$, NULL, $j${text}$j$::jsonb, {len(text)}, NOW())
ON CONFLICT (snapshot_date, section, endpoint) DO UPDATE SET
payload = EXCLUDED.payload,
payload_size = EXCLUDED.payload_size,
fetched_at = NOW();
"""
psql_stdin('', sql)
loaded += 1
by_section.setdefault(section, 0)
by_section[section] += 1
print(f'\nLOADED {loaded} files, SKIPPED {skipped}')
print('--- by section ---')
for s, n in sorted(by_section.items()):
print(f' {s:22s} {n}')
# Final summary from DB
print('\n--- DB summary ---')
out = psql(f"""
SELECT section, COUNT(*) AS endpoints, pg_size_pretty(SUM(payload_size)::bigint) AS total
FROM domrf_raw_endpoints
WHERE snapshot_date = '{SNAP}'
GROUP BY section
ORDER BY section;
""")
print(out)
if __name__ == '__main__':
main()