fs-lawrisk/lawrisk/utils/export_risk_json.py

76 lines
2.3 KiB
Python

import json
import os
import sys
import pg8000
from env_loader import load_env
# Read DB config from environment; provide sensible defaults
CONFIG = {
'host': os.getenv('PG_HOST', '8.138.196.105'),
'port': int(os.getenv('PG_PORT', '5432')),
'database': os.getenv('PG_DATABASE', 'fs_law_risk'),
'user': os.getenv('PG_USER', 'postgres'),
'password': os.getenv('PG_PASSWORD', 'difyai123456'),
}
# Export file path can be overridden via env
OUTPUT = os.getenv('RISK_EXPORT_OUTPUT', 'risk_tables_export.json')
SQLS = {
'risk_subject': "SELECT sub_id AS id, sub_name AS name FROM public.risk_subject ORDER BY sub_name;",
'risk_permit': "SELECT per_id AS id, per_name AS name FROM public.risk_permit ORDER BY per_name;",
'risk_sub_per': "SELECT sub_id, per_id FROM public.risk_sub_per;",
}
def fetch_all(cursor):
cols = [d[0] for d in cursor.description]
return [dict(zip(cols, row)) for row in cursor.fetchall()]
def main():
load_env()
try:
conn = pg8000.connect(**CONFIG)
except Exception as e:
print('DB connect error:', e, file=sys.stderr)
sys.exit(2)
result = {}
try:
with conn.cursor() as cur:
cur.execute(SQLS['risk_subject'])
subjects = fetch_all(cur)
cur.execute(SQLS['risk_permit'])
permits = fetch_all(cur)
cur.execute(SQLS['risk_sub_per'])
rels = fetch_all(cur)
conn.commit()
finally:
conn.close()
# Build mapping: subject -> list of permit_ids
sub_to_permit_ids = {}
for r in rels:
sub_id = r['sub_id']
per_id = r['per_id']
sub_to_permit_ids.setdefault(sub_id, set()).add(per_id)
# Subjects with aggregated permit_ids
subjects_out = []
for s in subjects:
subjects_out.append({
'id': s['id'],
'name': s['name'],
'permit_ids': sorted(list(sub_to_permit_ids.get(s['id'], [])))
})
# Final JSON: keep full permit catalog (id+name), and subjects contain aggregated permit_ids
result['risk_subject'] = subjects_out
result['risk_permit'] = permits
with open(OUTPUT, 'w', encoding='utf-8') as f:
json.dump(result, f, ensure_ascii=False, indent=2)
print('Exported to', OUTPUT)
if __name__ == '__main__':
main()