69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Reset database to a fresh state by dropping all tables and recreating them."""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
# Add project root to path
|
||
|
|
project_root = Path(__file__).parent.parent
|
||
|
|
if str(project_root) not in sys.path:
|
||
|
|
sys.path.insert(0, str(project_root))
|
||
|
|
|
||
|
|
from sqlalchemy import text
|
||
|
|
from src.core.database import Base, Database
|
||
|
|
from src.core.config import get_config
|
||
|
|
|
||
|
|
|
||
|
|
async def reset_database():
|
||
|
|
"""Reset database by dropping all tables and recreating them."""
|
||
|
|
config = get_config()
|
||
|
|
db_type = config.get("database.type", "postgresql")
|
||
|
|
|
||
|
|
print(f"Resetting {db_type} database...")
|
||
|
|
|
||
|
|
# Create database instance to get engine
|
||
|
|
db = Database()
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Drop all tables
|
||
|
|
# For PostgreSQL, we need to handle foreign key constraints
|
||
|
|
print("Dropping all tables...")
|
||
|
|
async with db.engine.begin() as conn:
|
||
|
|
# Disable multiple statements in one call caution, but reset logic usually needs strict control.
|
||
|
|
# However, asyncpg doesn't support "SET session_replication_role" easily within a transaction block
|
||
|
|
# if it's not a superuser or specific config.
|
||
|
|
# Instead, we will use CASCADE drop which is cleaner for full reset.
|
||
|
|
|
||
|
|
# Use reflection to find tables is harder in async.
|
||
|
|
# We will just drop all tables using CASCADE via raw SQL or metadata.
|
||
|
|
|
||
|
|
# Since we are using async engine, we need to utilize run_sync for metadata operations
|
||
|
|
# But drop_all doesn't support CASCADE automatically for all dialects in the way we might want
|
||
|
|
# if using pure SQLAlchemy metadata.drop_all.
|
||
|
|
|
||
|
|
# Let's try the standard metadata.drop_all first
|
||
|
|
await conn.run_sync(Base.metadata.drop_all)
|
||
|
|
|
||
|
|
print("Dropped all tables.")
|
||
|
|
|
||
|
|
# Recreate all tables
|
||
|
|
print("Recreating all tables...")
|
||
|
|
await db.create_tables()
|
||
|
|
print("Database reset complete!")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error during reset: {e}")
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
# Close database connection
|
||
|
|
await db.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
try:
|
||
|
|
asyncio.run(reset_database())
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error resetting database: {e}", file=sys.stderr)
|
||
|
|
sys.exit(1)
|