cleanup.py

import os
import psycopg2

def main():
    dsn = "postgresql://dira:dira_password@localhost:5433/dira_db"
    env_path = "./.env"
    if os.path.exists(env_path):
        with open(env_path, "r") as f:
            for line in f:
                if line.startswith("DATABASE_URL="):
                    dsn = line.split("=", 1)[1].strip().strip('"').strip("'")
                    break

    print(f"Connecting to database: {dsn}")
    
    conn = psycopg2.connect(dsn)
    conn.autocommit = True
    
    with conn.cursor() as cur:
        # Delete existing vendors with the phone number '+254712345678'
        cur.execute("DELETE FROM vendors WHERE phone_number = %s", ('+254712345678',))
        deleted_vendors = cur.rowcount
        print(f"Deleted {deleted_vendors} existing vendor records.")
        
        # Delete existing auth_users with temporary email addresses
        cur.execute("DELETE FROM auth_users WHERE email LIKE 'mama_%'")
        deleted_users = cur.rowcount
        print(f"Deleted {deleted_users} existing temporary auth_users records.")
        
        # Clean up any residual supply gaps for the items we will test
        cur.execute("DELETE FROM supply_gaps WHERE item_name IN ('Mango Juice', 'Maji Baridi')")
        deleted_gaps = cur.rowcount
        print(f"Deleted {deleted_gaps} residual supply gap records.")

    conn.close()
    print("Database cleanup completed successfully.")

if __name__ == "__main__":
    main()