Spaces:
Runtime error
Runtime error
| import asyncio | |
| from sqlalchemy import text | |
| from database.database import db_manager | |
| from logger.logger import get_logger | |
| from database.database_models import Base | |
| logger = get_logger(__name__) | |
| async def clean_all_data(): | |
| """ | |
| Clean all data from tables without dropping schema | |
| """ | |
| await db_manager.initialize() | |
| async with db_manager.engine.begin() as conn: | |
| logger.info("Cleaning all tables...") | |
| # Truncate tables in correct order to handle foreign key constraints | |
| # Using CASCADE will automatically handle dependent tables | |
| await conn.execute(text("TRUNCATE TABLE messages CASCADE;")) | |
| await conn.execute(text("TRUNCATE TABLE conversations CASCADE;")) | |
| await conn.execute(text("TRUNCATE TABLE documents CASCADE;")) | |
| await conn.execute(text("TRUNCATE TABLE users CASCADE;")) | |
| # Reset sequences | |
| await conn.execute(text(""" | |
| SELECT setval(pg_get_serial_sequence('users', 'id'), | |
| COALESCE((SELECT MAX(id) FROM users), 1), false); | |
| """)) | |
| await conn.execute(text(""" | |
| SELECT setval(pg_get_serial_sequence('documents', 'id'), | |
| COALESCE((SELECT MAX(id) FROM documents), 1), false); | |
| """)) | |
| logger.info("All tables cleaned and sequences reset!") | |
| async def drop_all_tables(): | |
| """ | |
| Drop all database tables. | |
| """ | |
| await db_manager.initialize() | |
| async with db_manager.engine.begin() as conn: | |
| logger.info("Dropping all tables...") | |
| await conn.run_sync(Base.metadata.drop_all) | |
| logger.info("All tables dropped successfully!") | |
| async def create_all_tables(): | |
| await db_manager.initialize() | |
| async with db_manager.engine.begin() as conn: | |
| logger.info("Creating tables...") | |
| await conn.run_sync(Base.metadata.create_all) | |
| logger.info("All tables created successfully!") | |
| if __name__ == "__main__": | |
| asyncio.run(create_all_tables()) |