init_db.py

"""Database initialization script"""
import asyncio
from database import init_db, engine
from db_models import Base
from rich.console import Console

console = Console()


async def create_tables():
    """Create all database tables"""
    console.print("\n[bold cyan]Initializing Trinka Database...[/bold cyan]")

    try:
        # Create all tables
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        console.print("[bold green]✓ Database tables created successfully![/bold green]")
        console.print("\n[bold]Tables created:[/bold]")
        console.print("  • conversations")
        console.print("  • messages")
        console.print("  • agent_actions")
        console.print("  • alert_groups")
        console.print("  • users (optional)")
        console.print("  • documents (optional)")

    except Exception as e:
        console.print(f"[bold red]✗ Error creating tables: {e}[/bold red]")
        raise

    finally:
        await engine.dispose()


async def drop_tables():
    """Drop all database tables (use with caution!)"""
    console.print("\n[bold yellow]⚠ Dropping all database tables...[/bold yellow]")

    try:
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.drop_all)

        console.print("[bold green]✓ All tables dropped successfully![/bold green]")

    except Exception as e:
        console.print(f"[bold red]✗ Error dropping tables: {e}[/bold red]")
        raise

    finally:
        await engine.dispose()


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "--drop":
        console.print("[bold red]WARNING: This will delete all data![/bold red]")
        response = input("Are you sure you want to drop all tables? (yes/no): ")
        if response.lower() == "yes":
            asyncio.run(drop_tables())
        else:
            console.print("[yellow]Operation cancelled.[/yellow]")
    else:
        asyncio.run(create_tables())