95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
AI Analysis Service Database Migration Runner
|
|
Runs the database migration for AI Analysis Service during container startup.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
|
|
def log(message):
|
|
"""Log with timestamp."""
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}")
|
|
|
|
def check_database_connection():
|
|
"""Check if database is available."""
|
|
try:
|
|
import psycopg2
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
conn = psycopg2.connect(
|
|
host=os.getenv('POSTGRES_HOST', 'localhost'),
|
|
port=os.getenv('POSTGRES_PORT', 5432),
|
|
database=os.getenv('POSTGRES_DB', 'dev_pipeline'),
|
|
user=os.getenv('POSTGRES_USER', 'pipeline_admin'),
|
|
password=os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
|
|
)
|
|
conn.close()
|
|
return True
|
|
except Exception as e:
|
|
log(f"Database connection failed: {e}")
|
|
return False
|
|
|
|
def run_migration():
|
|
"""Run the database migration."""
|
|
try:
|
|
log("Starting AI Analysis Service database migration...")
|
|
|
|
# Check if database is available
|
|
max_retries = 30
|
|
retry_count = 0
|
|
|
|
while retry_count < max_retries:
|
|
if check_database_connection():
|
|
log("Database connection successful")
|
|
break
|
|
else:
|
|
retry_count += 1
|
|
log(f"Database not ready, retrying in 2 seconds... ({retry_count}/{max_retries})")
|
|
time.sleep(2)
|
|
else:
|
|
log("ERROR: Could not connect to database after 60 seconds")
|
|
return False
|
|
|
|
# Run the migration script
|
|
schema_file = Path(__file__).parent / "001-schema.sql"
|
|
if not schema_file.exists():
|
|
log("ERROR: Schema file not found")
|
|
return False
|
|
|
|
log(f"Running migration from {schema_file}")
|
|
|
|
# Use psql to run the migration
|
|
env = os.environ.copy()
|
|
env['PGPASSWORD'] = os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
|
|
|
|
result = subprocess.run([
|
|
'psql',
|
|
'-h', os.getenv('POSTGRES_HOST', 'localhost'),
|
|
'-p', os.getenv('POSTGRES_PORT', '5432'),
|
|
'-U', os.getenv('POSTGRES_USER', 'pipeline_admin'),
|
|
'-d', os.getenv('POSTGRES_DB', 'dev_pipeline'),
|
|
'-f', str(schema_file),
|
|
'-v', 'ON_ERROR_STOP=1'
|
|
], env=env, capture_output=True, text=True)
|
|
|
|
if result.returncode == 0:
|
|
log("✅ AI Analysis Service database migration completed successfully")
|
|
return True
|
|
else:
|
|
log(f"❌ Migration failed: {result.stderr}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
log(f"❌ Migration error: {e}")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
success = run_migration()
|
|
sys.exit(0 if success else 1)
|