#!/usr/bin/env python3 """ Database initialization script for TimeTracker This script checks if the database is connected and initialized, and initializes it if needed. """ import os import sys import time import traceback from sqlalchemy import create_engine, text, inspect from sqlalchemy.exc import OperationalError, ProgrammingError def wait_for_database(url, max_attempts=30, delay=2): """Wait for database to be ready""" print(f"Waiting for database to be ready...") for attempt in range(max_attempts): try: engine = create_engine(url, pool_pre_ping=True) with engine.connect() as conn: conn.execute(text("SELECT 1")) print("Database connection established successfully") return engine except Exception as e: print(f"Waiting for database... (attempt {attempt+1}/{max_attempts}): {e}") if attempt < max_attempts - 1: time.sleep(delay) else: print("Database not ready after waiting, exiting...") sys.exit(1) return None def check_database_initialization(engine): """Check if database is initialized by looking for required tables and correct schema""" print("Checking if database is initialized...") try: inspector = inspect(engine) # Check if our main tables exist existing_tables = inspector.get_table_names() required_tables = ['users', 'projects', 'time_entries', 'settings'] missing_tables = [table for table in required_tables if table not in existing_tables] if missing_tables: print(f"Database not fully initialized. Missing tables: {missing_tables}") return False else: print("✓ All required tables exist") # Check if tables have the correct schema print("Checking table schemas...") # Check if time_entries has task_id column if 'time_entries' in existing_tables: time_entries_columns = [col['name'] for col in inspector.get_columns('time_entries')] print(f"Debug: time_entries columns found: {time_entries_columns}") if 'task_id' not in time_entries_columns: print(f"✗ time_entries table missing task_id column. Available columns: {time_entries_columns}") return False else: print("✓ time_entries table has correct schema") # Check if tasks table exists if 'tasks' not in existing_tables: print("⚠ tasks table missing - will be created by SQL script") # Don't return False here, let the SQL script handle it else: print("✓ tasks table exists") print("✓ Database is already initialized with all required tables and correct schema") return True except Exception as e: print(f"Error checking database initialization: {e}") print(f"Traceback: {traceback.format_exc()}") return False def check_table_schema(engine, table_name, required_columns): """Check if a table has the required columns""" try: inspector = inspect(engine) if table_name not in inspector.get_table_names(): return False existing_columns = [col['name'] for col in inspector.get_columns(table_name)] missing_columns = [col for col in required_columns if col not in existing_columns] if missing_columns: print(f"Table {table_name} missing columns: {missing_columns}") return False return True except Exception as e: print(f"Error checking schema for {table_name}: {e}") return False def ensure_correct_schema(engine): """Ensure all tables have the correct schema""" print("Checking table schemas...") # Define required columns for each table required_columns = { 'time_entries': ['id', 'user_id', 'project_id', 'task_id', 'start_time', 'end_time', 'duration_seconds', 'notes', 'tags', 'source', 'billable', 'created_at', 'updated_at'] # Note: tasks table is created by SQL script, not checked here } needs_recreation = False for table_name, columns in required_columns.items(): if not check_table_schema(engine, table_name, columns): print(f"Table {table_name} needs recreation") needs_recreation = True return needs_recreation def initialize_database(engine): """Initialize database using Flask CLI command""" print("Initializing database...") try: # Set environment variables for Flask os.environ['FLASK_APP'] = 'app' os.environ['FLASK_ENV'] = 'production' print("Importing Flask app...") # Import Flask app and initialize database from app import create_app, db from app.models import User, Project, TimeEntry, Settings print("Creating Flask app...") app = create_app() print("Setting up app context...") with app.app_context(): # Check if we need to recreate tables due to schema mismatch if ensure_correct_schema(engine): print("Schema mismatch detected, dropping and recreating tables...") db.drop_all() print("All tables dropped") print("Creating all tables...") # Create all tables db.create_all() print("Verifying tables were created...") # Verify tables were created inspector = inspect(engine) existing_tables = inspector.get_table_names() print(f"Tables after creation: {existing_tables}") # Create default admin user if it doesn't exist admin_username = os.getenv('ADMIN_USERNAMES', 'admin').split(',')[0] print(f"Checking for admin user: {admin_username}") if not User.query.filter_by(username=admin_username).first(): print("Creating admin user...") admin_user = User( username=admin_username, role='admin' ) admin_user.is_active = True db.session.add(admin_user) db.session.commit() print(f"Created default admin user: {admin_username}") else: print(f"Admin user {admin_username} already exists") # Create default settings if they don't exist print("Checking for default settings...") if not Settings.query.first(): print("Creating default settings...") settings = Settings() db.session.add(settings) db.session.commit() print("Created default settings") else: print("Default settings already exist") # Create default project if it doesn't exist print("Checking for default project...") if not Project.query.first(): print("Creating default project...") project = Project( name='General', client='Default Client', description='Default project for general tasks', billable=True, status='active' ) db.session.add(project) db.session.commit() print("Created default project") else: print("Default project already exists") print("Database initialized successfully") return True except Exception as e: print(f"Error initializing database: {e}") print(f"Traceback: {traceback.format_exc()}") return False def main(): """Main function""" url = os.getenv("DATABASE_URL", "") if not url.startswith("postgresql"): print("No PostgreSQL database configured, skipping initialization") return print(f"Database URL: {url}") # Wait for database to be ready engine = wait_for_database(url) # Check if database is initialized print("=== Starting database initialization check ===") if not check_database_initialization(engine): print("=== Database not initialized, starting initialization ===") # Initialize database if initialize_database(engine): print("Database initialization completed successfully") # Verify initialization worked if check_database_initialization(engine): print("Database verification successful") else: print("Database verification failed - tables still missing") sys.exit(1) else: print("Database initialization failed") sys.exit(1) else: print("=== Database already initialized, checking if reinitialization is needed ===") # Even if database is initialized, double-check schema and reinitialize if needed print("Double-checking schema for existing database...") if ensure_correct_schema(engine): print("Schema mismatch detected in existing database, reinitializing...") if initialize_database(engine): print("Database reinitialization completed successfully") # Verify reinitialization worked if check_database_initialization(engine): print("Database verification successful after reinitialization") else: print("Database verification failed after reinitialization") sys.exit(1) else: print("Database reinitialization failed") sys.exit(1) else: print("Schema is correct, no reinitialization needed") if __name__ == "__main__": main()