#!/usr/bin/env python
"""
Script to drop all tables in the Preloop database.
Use with caution - this will delete all data!
"""
import os
import sys
import click
import psycopg2
from dotenv import load_dotenv
@click.command()
@click.option("--force", is_flag=True, help="Skip confirmation and force deletion")
def drop_tables(force: bool):
"""
Drop all tables in the Preloop database.
"""
# Load environment variables
load_dotenv()
# Get database connection string from environment
db_url = os.getenv("DATABASE_URL")
if not db_url:
click.echo("ERROR: DATABASE_URL environment variable not set.")
sys.exit(1)
# Strip SQLAlchemy dialect suffix (e.g. postgresql+psycopg:// -> postgresql://)
db_url = db_url.replace("postgresql+psycopg://", "postgresql://")
# Connect to database
try:
conn = psycopg2.connect(db_url)
conn.autocommit = True
cursor = conn.cursor()
except Exception as e:
click.echo(f"ERROR: Could not connect to database: {str(e)}")
sys.exit(1)
# Get list of tables
try:
cursor.execute(
"""
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public'
ORDER BY table_name
"""
)
tables = [row[0] for row in cursor.fetchall()]
except Exception as e:
click.echo(f"ERROR: Could not get table list: {str(e)}")
conn.close()
sys.exit(1)
if not tables:
click.echo("No tables found in database.")
conn.close()
return
if not force:
click.echo("WARNING: This will delete ALL data in the database.")
click.echo("Tables to be dropped:")
for table in tables:
click.echo(f" - {table}")
# Ask for confirmation
if not click.confirm("Are you sure you want to continue?"):
click.echo("Operation cancelled.")
conn.close()
sys.exit(0)
click.echo("Dropping all tables...")
# Drop all tables (disable foreign key constraints first)
try:
cursor.execute("SET session_replication_role = 'replica';")
for table in tables:
cursor.execute(f'DROP TABLE IF EXISTS "{table}" CASCADE;')
cursor.execute("SET session_replication_role = 'origin';")
click.echo("All tables have been dropped successfully.")
except Exception as e:
click.echo(f"ERROR: Could not drop tables: {str(e)}")
conn.close()
return
# Drop all custom enum types
try:
cursor.execute(
"""
SELECT typname FROM pg_type
JOIN pg_namespace ON pg_type.typnamespace = pg_namespace.oid
WHERE pg_namespace.nspname = 'public' AND pg_type.typtype = 'e'
"""
)
enums = [row[0] for row in cursor.fetchall()]
if enums:
for enum in enums:
cursor.execute(f'DROP TYPE IF EXISTS "{enum}" CASCADE;')
click.echo(f"Dropped {len(enums)} enum type(s).")
else:
click.echo("No enum types to drop.")
except Exception as e:
click.echo(f"WARNING: Could not drop enum types: {str(e)}")
finally:
conn.close()
if __name__ == "__main__":
drop_tables()
