import os import sys import subprocess import datetime import pg8000.dbapi as pg # ----------------------------------------------------------------------------- # 1. Environment & Configuration # ----------------------------------------------------------------------------- def load_env(): env_vars = {} try: with open('.env', 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line or line.startswith('#'): continue if '=' in line: k, v = line.split('=', 1) env_vars[k.strip()] = v.strip() except Exception: pass return env_vars env = load_env() HOST = env.get("LIC_PG_HOST", "172.24.240.1") PORT = int(env.get("LIC_PG_PORT", "5432")) USER = env.get("LIC_PG_USER", "postgres") PASSWORD = env.get("LIC_PG_PASSWORD", "") DATABASE = env.get("LIC_PG_DATABASE", "licensing_risks") # ----------------------------------------------------------------------------- # 2. Database Backup # ----------------------------------------------------------------------------- def perform_backup(): date_str = datetime.datetime.now().strftime("%Y-%m-%d") backup_file = os.path.abspath(f"data/backup_{date_str}.sql") # Ensure data directory exists os.makedirs(os.path.dirname(backup_file), exist_ok=True) print(f"Starting backup to {backup_file}...") # Set PGPASSWORD for pg_dump env_copy = os.environ.copy() env_copy["PGPASSWORD"] = PASSWORD # Construct command # pg_dump -h host -p port -U user -f file dbname cmd = [ "pg_dump", "-h", HOST, "-p", str(PORT), "-U", USER, "-f", backup_file, DATABASE ] try: subprocess.run(cmd, env=env_copy, check=True) print("Backup completed successfully.") return True except subprocess.CalledProcessError as e: print(f"Backup failed: {e}") return False except FileNotFoundError: print("pg_dump not found in PATH. Skipping backup (RISKY!).") # In this specific task, if backup fails, we should probably stop or ask user. # But given I am an agent, I should try to warn and maybe proceed if trivial, but deleting data is dangerous. # I will stop if backup fails. return False # ----------------------------------------------------------------------------- # 3. Data Cleanup # ----------------------------------------------------------------------------- def perform_cleanup(): print("Connecting to database...") try: conn = pg.connect(host=HOST, port=PORT, user=USER, password=PASSWORD, database=DATABASE) cursor = conn.cursor() target_theme_name = "不涉及" # 1. Find theme ID print(f"Finding theme '{target_theme_name}'...") cursor.execute("SELECT id FROM themes WHERE name = %s", (target_theme_name,)) row = cursor.fetchone() if not row: print(f"Theme '{target_theme_name}' not found. Nothing to do.") conn.close() return theme_id = row[0] print(f"Found theme_id: {theme_id}") # 2. Find associated permits (matters) # Using region_theme_permits as the binding table print("Finding associated permits...") cursor.execute("SELECT DISTINCT permit_id FROM region_theme_permits WHERE theme_id = %s", (theme_id,)) permit_rows = cursor.fetchall() permit_ids = [r[0] for r in permit_rows] print(f"Found {len(permit_ids)} permits bound to '{target_theme_name}'.") if not permit_ids: print("No permits bound to this theme. Proceeding to delete theme only.") # Start Transaction try: # 3. Clear TOPICS for these permits (Unbind ALL themes from these permits) if permit_ids: # Need to convert UUIDs to string for SQL IN clause or use execute with any # pg8000 handles list/tuple for IN nicely if formatted manually, or we can loop. # Batch delete is better. # DELETE FROM region_theme_permits WHERE permit_id IN (...) # Format for IN clause permit_ids_tuple = tuple(permit_ids) if len(permit_ids) == 1: # distinct check ensures > 0, but tuple of 1 element needs comma pass # Note: pg8000 might struggle with large IN clauses in param logic depending on version. # But let's try standard parameterized query. # Actually, to be safe with syntax: "IN %s" with a tuple works in some drivers, # but often 'IN (%s, %s)' is needed. # I'll generate placeholders. placeholders = ",".join(["%s"] * len(permit_ids)) print(f"Clearing ALL topic bindings for these {len(permit_ids)} permits...") # Delete from region_theme_permits cursor.execute(f"DELETE FROM region_theme_permits WHERE permit_id IN ({placeholders})", permit_ids_tuple) print(f"Deleted {cursor.rowcount} rows from region_theme_permits.") # Also check region_permit_theme_overrides if they exist for these permits? # The user said "clear topics". Overrides might not be "bindings" per se, but let's check table existence. # I'll skip overrides to avoid complexity unless requested, sticking to "bindings". # 4. Delete the theme itself from all related tables print(f"Deleting theme '{target_theme_name}' references...") # Delete from region_themes cursor.execute("DELETE FROM region_themes WHERE theme_id = %s", (theme_id,)) print(f"Deleted {cursor.rowcount} rows from region_themes.") # Delete from permit_theme_rules cursor.execute("DELETE FROM permit_theme_rules WHERE theme_id = %s", (theme_id,)) print(f"Deleted {cursor.rowcount} rows from permit_theme_rules.") # Delete from themes cursor.execute("DELETE FROM themes WHERE id = %s", (theme_id,)) print(f"Deleted {cursor.rowcount} row from themes.") # Commit conn.commit() print("Transaction committed successfully.") except Exception as e: conn.rollback() print(f"Error during transaction, rolled back. Error: {e}") raise conn.close() except Exception as e: print(f"Database error: {e}") if __name__ == "__main__": if perform_backup(): perform_cleanup() else: print("Aborting cleanup because backup failed.")