Files
TimeTracker/docker/startup_with_migration.py
Dries Peeters 3894b3be94 Migration
2025-09-02 15:42:32 +02:00

395 lines
14 KiB
Python

#!/usr/bin/env python3
"""
Enhanced Startup Script with Automatic Migration Detection
This script automatically detects database state and chooses the correct migration strategy
"""
import os
import sys
import time
import subprocess
import sqlite3
import psycopg2
from pathlib import Path
from datetime import datetime
import logging
# Configure logging
try:
# Ensure logs directory exists
os.makedirs('/app/logs', exist_ok=True)
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s',
handlers=[
logging.StreamHandler(sys.stdout),
logging.FileHandler('/app/logs/timetracker_startup.log')
]
)
except Exception as e:
# Fallback to console-only logging if file logging fails
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s',
handlers=[logging.StreamHandler(sys.stdout)]
)
print(f"Warning: Could not set up file logging: {e}")
logger = logging.getLogger(__name__)
def wait_for_database(db_url, max_retries=60, retry_delay=3):
"""Wait for database to be available"""
logger.info(f"Waiting for database to be available: {db_url}")
for attempt in range(max_retries):
try:
if db_url.startswith('postgresql'):
# Handle both postgresql:// and postgresql+psycopg2:// URLs
clean_url = db_url.replace('+psycopg2://', '://')
conn = psycopg2.connect(clean_url)
conn.close()
logger.info("✓ PostgreSQL database is available")
return True
elif db_url.startswith('sqlite'):
# For SQLite, just check if the file exists or can be created
db_file = db_url.replace('sqlite:///', '')
if os.path.exists(db_file) or os.access(os.path.dirname(db_file), os.W_OK):
logger.info("✓ SQLite database is available")
return True
except Exception as e:
logger.info(f"Database not ready (attempt {attempt + 1}/{max_retries}): {e}")
if attempt < max_retries - 1:
time.sleep(retry_delay)
logger.error("Database is not available after maximum retries")
return False
def detect_database_state(db_url):
"""Detect the current state of the database"""
logger.info("Analyzing database state...")
try:
if db_url.startswith('postgresql'):
# Handle both postgresql:// and postgresql+psycopg2:// URLs
clean_url = db_url.replace('+psycopg2://', '://')
conn = psycopg2.connect(clean_url)
cursor = conn.cursor()
# Check if alembic_version table exists
cursor.execute("""
SELECT EXISTS (
SELECT FROM information_schema.tables
WHERE table_name = 'alembic_version'
)
""")
has_alembic = cursor.fetchone()[0]
# Get list of existing tables
cursor.execute("""
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public'
ORDER BY table_name
""")
existing_tables = [row[0] for row in cursor.fetchall()]
# Check if this is a fresh database
is_fresh = len(existing_tables) == 0
conn.close()
logger.info(f"Database state: has_alembic={has_alembic}, tables={existing_tables}, is_fresh={is_fresh}")
if has_alembic:
return 'migrated', existing_tables
elif existing_tables:
return 'legacy', existing_tables
else:
return 'fresh', []
elif db_url.startswith('sqlite'):
db_file = db_url.replace('sqlite:///', '')
if not os.path.exists(db_file):
return 'fresh', []
conn = sqlite3.connect(db_file)
cursor = conn.cursor()
# Check if alembic_version table exists
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='alembic_version'")
has_alembic = cursor.fetchone() is not None
# Get list of existing tables
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
existing_tables = [row[0] for row in cursor.fetchall()]
conn.close()
logger.info(f"Database state: has_alembic={has_alembic}, tables={existing_tables}")
if has_alembic:
return 'migrated', existing_tables
elif existing_tables:
return 'legacy', existing_tables
else:
return 'fresh', []
except Exception as e:
logger.error(f"Error detecting database state: {e}")
return 'unknown', []
return 'unknown', []
def choose_migration_strategy(db_state, existing_tables):
"""Choose the appropriate migration strategy based on database state"""
logger.info(f"Choosing migration strategy for state: {db_state}")
if db_state == 'fresh':
logger.info("Fresh database detected - using standard initialization")
return 'fresh_init'
elif db_state == 'migrated':
logger.info("Database already migrated - checking for pending migrations")
return 'check_migrations'
elif db_state == 'legacy':
logger.info("Legacy database detected - using comprehensive migration")
return 'comprehensive_migration'
else:
logger.warning("Unknown database state - using comprehensive migration as fallback")
return 'comprehensive_migration'
def execute_migration_strategy(strategy, db_url):
"""Execute the chosen migration strategy"""
logger.info(f"Executing migration strategy: {strategy}")
try:
if strategy == 'fresh_init':
return execute_fresh_init(db_url)
elif strategy == 'check_migrations':
return execute_check_migrations(db_url)
elif strategy == 'comprehensive_migration':
return execute_comprehensive_migration(db_url)
else:
logger.error(f"Unknown migration strategy: {strategy}")
return False
except Exception as e:
logger.error(f"Error executing migration strategy: {e}")
return False
def execute_fresh_init(db_url):
"""Execute fresh database initialization"""
logger.info("Executing fresh database initialization...")
try:
# Initialize Flask-Migrate
result = subprocess.run(['flask', 'db', 'init'],
capture_output=True, text=True, check=True)
logger.info("✓ Flask-Migrate initialized")
# Create initial migration
result = subprocess.run(['flask', 'db', 'migrate', '-m', 'Initial database schema'],
capture_output=True, text=True, check=True)
logger.info("✓ Initial migration created")
# Apply migration
result = subprocess.run(['flask', 'db', 'upgrade'],
capture_output=True, text=True, check=True)
logger.info("✓ Initial migration applied")
return True
except subprocess.CalledProcessError as e:
logger.error(f"Fresh init failed: {e}")
logger.error(f"STDOUT: {e.stdout}")
logger.error(f"STDERR: {e.stderr}")
return False
def execute_check_migrations(db_url):
"""Check and apply any pending migrations"""
logger.info("Checking for pending migrations...")
try:
# Check current migration status
result = subprocess.run(['flask', 'db', 'current'],
capture_output=True, text=True, check=True)
current_revision = result.stdout.strip()
logger.info(f"Current migration revision: {current_revision}")
# Check for pending migrations
result = subprocess.run(['flask', 'db', 'upgrade'],
capture_output=True, text=True, check=True)
logger.info("✓ Migrations checked and applied")
return True
except subprocess.CalledProcessError as e:
logger.error(f"Migration check failed: {e}")
return False
def execute_comprehensive_migration(db_url):
"""Execute comprehensive migration for legacy databases"""
logger.info("Executing comprehensive migration...")
try:
# Run the comprehensive migration script
migration_script = '/app/migrations/migrate_existing_database.py'
if os.path.exists(migration_script):
result = subprocess.run(['python', migration_script],
capture_output=True, text=True, check=True)
logger.info("✓ Comprehensive migration completed")
return True
else:
logger.warning("Comprehensive migration script not found, falling back to manual migration")
return execute_manual_migration(db_url)
except subprocess.CalledProcessError as e:
logger.error(f"Comprehensive migration failed: {e}")
logger.error(f"STDOUT: {e.stdout}")
logger.error(f"STDERR: {e.stderr}")
return False
def execute_manual_migration(db_url):
"""Execute manual migration as fallback"""
logger.info("Executing manual migration fallback...")
try:
# Initialize Flask-Migrate if not already done
if not os.path.exists('/app/migrations/env.py'):
result = subprocess.run(['flask', 'db', 'init'],
capture_output=True, text=True, check=True)
logger.info("✓ Flask-Migrate initialized")
# Create baseline migration
result = subprocess.run(['flask', 'db', 'migrate', '-m', 'Baseline from existing database'],
capture_output=True, text=True, check=True)
logger.info("✓ Baseline migration created")
# Stamp database as current
result = subprocess.run(['flask', 'db', 'stamp', 'head'],
capture_output=True, text=True, check=True)
logger.info("✓ Database stamped as current")
return True
except subprocess.CalledProcessError as e:
logger.error(f"Manual migration failed: {e}")
return False
def verify_database_integrity(db_url):
"""Verify that the database is working correctly after migration"""
logger.info("Verifying database integrity...")
try:
# Test basic database operations
if db_url.startswith('postgresql'):
# Handle both postgresql:// and postgresql+psycopg2:// URLs
clean_url = db_url.replace('+psycopg2://', '://')
conn = psycopg2.connect(clean_url)
cursor = conn.cursor()
# Check if key tables exist
cursor.execute("""
SELECT table_name
FROM information_schema.tables
WHERE table_name IN ('users', 'projects', 'time_entries')
AND table_schema = 'public'
""")
key_tables = [row[0] for row in cursor.fetchall()]
conn.close()
if len(key_tables) >= 2: # At least users and projects
logger.info("✓ Database integrity verified")
return True
else:
logger.error(f"Missing key tables: {key_tables}")
return False
elif db_url.startswith('sqlite'):
db_file = db_url.replace('sqlite:///', '')
if not os.path.exists(db_file):
logger.error("SQLite database file not found")
return False
conn = sqlite3.connect(db_file)
cursor = conn.cursor()
# Check if key tables exist
cursor.execute("""
SELECT name FROM sqlite_master
WHERE type='table' AND name IN ('users', 'projects', 'time_entries')
""")
key_tables = [row[0] for row in cursor.fetchall()]
conn.close()
if len(key_tables) >= 2: # At least users and projects
logger.info("✓ Database integrity verified")
return True
else:
logger.error(f"Missing key tables: {key_tables}")
return False
except Exception as e:
logger.error(f"Database integrity check failed: {e}")
return False
return False
def main():
"""Main startup function"""
logger.info("=== TimeTracker Enhanced Startup with Migration Detection ===")
# Set environment variables
os.environ.setdefault('FLASK_APP', '/app/app.py')
# Get database URL from environment
db_url = os.getenv('DATABASE_URL')
if not db_url:
logger.error("DATABASE_URL environment variable not set")
sys.exit(1)
logger.info(f"Database URL: {db_url}")
# Wait for database to be available
if not wait_for_database(db_url):
logger.error("Failed to connect to database")
sys.exit(1)
# Detect database state
db_state, existing_tables = detect_database_state(db_url)
logger.info(f"Detected database state: {db_state} with {len(existing_tables)} tables")
# Choose migration strategy
strategy = choose_migration_strategy(db_state, existing_tables)
logger.info(f"Selected migration strategy: {strategy}")
# Execute migration strategy
if not execute_migration_strategy(strategy, db_url):
logger.error("Migration strategy execution failed")
sys.exit(1)
# Verify database integrity
if not verify_database_integrity(db_url):
logger.error("Database integrity verification failed")
sys.exit(1)
logger.info("=== Startup and Migration Complete ===")
logger.info("Database is ready for use")
# Show final migration status
try:
result = subprocess.run(['flask', 'db', 'current'],
capture_output=True, text=True, check=True)
logger.info(f"Final migration status: {result.stdout.strip()}")
except:
logger.info("Could not determine final migration status")
if __name__ == "__main__":
main()