116 lines
3.0 KiB
Python
116 lines
3.0 KiB
Python
"""
|
|
Database initialization script
|
|
|
|
This script creates all database tables based on the defined models.
|
|
Run this script to set up the database schema for the first time.
|
|
|
|
Usage:
|
|
python scripts/init_db.py
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add the project root to the Python path
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
from src.database.models import Base
|
|
from src.database.session import db_manager
|
|
from src.core.config import get_config
|
|
|
|
|
|
def init_database():
|
|
"""Initialize the database by creating all tables"""
|
|
print("Initializing database...")
|
|
|
|
# Load configuration
|
|
try:
|
|
config = get_config()
|
|
print(f"Database URL: {config.database.url}")
|
|
except Exception as e:
|
|
print(f"Error loading configuration: {e}")
|
|
sys.exit(1)
|
|
|
|
# Initialize database manager
|
|
try:
|
|
db_manager.initialize()
|
|
engine = db_manager.get_engine()
|
|
except Exception as e:
|
|
print(f"Error connecting to database: {e}")
|
|
sys.exit(1)
|
|
|
|
# Create all tables
|
|
try:
|
|
Base.metadata.create_all(bind=engine)
|
|
print("Database tables created successfully!")
|
|
|
|
# Display created tables
|
|
print("\nCreated tables:")
|
|
for table_name in Base.metadata.tables.keys():
|
|
print(f" - {table_name}")
|
|
|
|
except Exception as e:
|
|
print(f"Error creating database tables: {e}")
|
|
sys.exit(1)
|
|
finally:
|
|
db_manager.close()
|
|
|
|
print("\nDatabase initialization complete!")
|
|
|
|
|
|
def drop_database():
|
|
"""Drop all database tables (USE WITH CAUTION!)"""
|
|
print("WARNING: This will drop all tables and delete all data!")
|
|
response = input("Are you sure you want to continue? (yes/no): ")
|
|
|
|
if response.lower() != "yes":
|
|
print("Operation cancelled.")
|
|
return
|
|
|
|
print("Dropping all database tables...")
|
|
|
|
try:
|
|
config = get_config()
|
|
db_manager.initialize()
|
|
engine = db_manager.get_engine()
|
|
|
|
Base.metadata.drop_all(bind=engine)
|
|
print("All database tables dropped successfully!")
|
|
|
|
except Exception as e:
|
|
print(f"Error dropping database tables: {e}")
|
|
sys.exit(1)
|
|
finally:
|
|
db_manager.close()
|
|
|
|
|
|
def reset_database():
|
|
"""Drop and recreate all database tables"""
|
|
drop_database()
|
|
print("\nRecreating database...")
|
|
init_database()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Database initialization script")
|
|
parser.add_argument(
|
|
"command",
|
|
choices=["init", "drop", "reset"],
|
|
nargs="?",
|
|
default="init",
|
|
help="Command to execute (init: create tables, drop: remove tables, reset: drop and recreate)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.command == "init":
|
|
init_database()
|
|
elif args.command == "drop":
|
|
drop_database()
|
|
elif args.command == "reset":
|
|
reset_database()
|
|
|