fs-lawrisk/fix_unit_levels.py

#!/usr/bin/env python
"""Fix unit_level data based on code patterns"""
import sys
import os

# Add project root to path
project_root = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, project_root)

# Load .env
from lawrisk.utils.env_loader import load_env
load_env('.env', override=False)

# Import database connection
import pg8000

def main():
    print("=" * 60)
    print("Fixing unit_level data")
    print("=" * 60)
    print()

    # Connect to database
    conn = pg8000.connect(
        host=os.getenv('LIC_PG_HOST'),
        port=int(os.getenv('LIC_PG_PORT', '5432')),
        user=os.getenv('LIC_PG_USER'),
        password=os.getenv('LIC_PG_PASSWORD'),
        database=os.getenv('LIC_PG_DATABASE')
    )
    cur = conn.cursor()

    # Fix市级管理员 (FSSJSJ)
    print("Fixing FSSJSJ (city admin)...")
    cur.execute("""
        UPDATE service_departments
        SET unit_level = 'admin', grade = 90
        WHERE code = 'FSSJSJ'
    """)
    print(f"  Updated {cur.rowcount} row(s)")

    # Fix区局管理员 (parent_id IS NOT NULL means child departments)
    print("\nFixing district administrations (parent_id is not null)...")
    cur.execute("""
        UPDATE service_departments
        SET unit_level = 'district', grade = 80
        WHERE code != 'FSSJSJ'
    """)
    print(f"  Updated {cur.rowcount} row(s)")

    # Commit changes
    conn.commit()

    # Verify results
    print("\n" + "=" * 60)
    print("Verification")
    print("=" * 60)
    print()

    cur.execute("""
        SELECT code, name, unit_level, grade
        FROM service_departments
        ORDER BY code
    """)

    for code, name, unit_level, grade in cur.fetchall():
        print(f"{code}: {unit_level} (grade={grade})")

    conn.close()

    print("\n" + "=" * 60)
    print("Done! Unit levels have been corrected.")
    print("=" * 60)

if __name__ == "__main__":
    main()