136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Database initialization script for PostgreSQL.
|
|
Replaces psql commands in startup.sh with pure Python.
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
import psycopg2
|
|
from psycopg2 import OperationalError
|
|
|
|
DB_HOST = os.getenv("DB_HOST", "localhost")
|
|
DB_PORT = os.getenv("DB_PORT", "5432")
|
|
DB_NAME = os.getenv("DB_NAME", "mitai_dev")
|
|
DB_USER = os.getenv("DB_USER", "mitai_dev")
|
|
DB_PASSWORD = os.getenv("DB_PASSWORD", "")
|
|
|
|
def get_connection():
|
|
"""Get PostgreSQL connection."""
|
|
return psycopg2.connect(
|
|
host=DB_HOST,
|
|
port=DB_PORT,
|
|
database=DB_NAME,
|
|
user=DB_USER,
|
|
password=DB_PASSWORD
|
|
)
|
|
|
|
def wait_for_postgres(max_retries=30):
|
|
"""Wait for PostgreSQL to be ready."""
|
|
print("\nChecking PostgreSQL connection...")
|
|
for i in range(1, max_retries + 1):
|
|
try:
|
|
conn = get_connection()
|
|
conn.close()
|
|
print("✓ PostgreSQL ready")
|
|
return True
|
|
except OperationalError:
|
|
print(f" Waiting for PostgreSQL... (attempt {i}/{max_retries})")
|
|
time.sleep(2)
|
|
|
|
print(f"✗ PostgreSQL not ready after {max_retries} attempts")
|
|
return False
|
|
|
|
def check_table_exists(table_name="profiles"):
|
|
"""Check if a table exists."""
|
|
try:
|
|
conn = get_connection()
|
|
cur = conn.cursor()
|
|
cur.execute("""
|
|
SELECT COUNT(*)
|
|
FROM information_schema.tables
|
|
WHERE table_schema='public' AND table_name=%s
|
|
""", (table_name,))
|
|
count = cur.fetchone()[0]
|
|
cur.close()
|
|
conn.close()
|
|
return count > 0
|
|
except Exception as e:
|
|
print(f"Error checking table: {e}")
|
|
return False
|
|
|
|
def load_schema(schema_file="/app/schema.sql"):
|
|
"""Load schema from SQL file."""
|
|
try:
|
|
with open(schema_file, 'r') as f:
|
|
schema_sql = f.read()
|
|
|
|
conn = get_connection()
|
|
cur = conn.cursor()
|
|
cur.execute(schema_sql)
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
print("✓ Schema loaded from schema.sql")
|
|
return True
|
|
except Exception as e:
|
|
print(f"✗ Error loading schema: {e}")
|
|
return False
|
|
|
|
def get_profile_count():
|
|
"""Get number of profiles in database."""
|
|
try:
|
|
conn = get_connection()
|
|
cur = conn.cursor()
|
|
cur.execute("SELECT COUNT(*) FROM profiles")
|
|
count = cur.fetchone()[0]
|
|
cur.close()
|
|
conn.close()
|
|
return count
|
|
except Exception as e:
|
|
print(f"Error getting profile count: {e}")
|
|
return -1
|
|
|
|
if __name__ == "__main__":
|
|
print("═══════════════════════════════════════════════════════════")
|
|
print("MITAI JINKENDO - Database Initialization (v9c)")
|
|
print("═══════════════════════════════════════════════════════════")
|
|
|
|
# Wait for PostgreSQL
|
|
if not wait_for_postgres():
|
|
sys.exit(1)
|
|
|
|
# Check schema
|
|
print("\nChecking database schema...")
|
|
if not check_table_exists("profiles"):
|
|
print(" Schema not found, initializing...")
|
|
if not load_schema():
|
|
sys.exit(1)
|
|
else:
|
|
print("✓ Schema already exists")
|
|
|
|
# Check for migration
|
|
print("\nChecking for SQLite data migration...")
|
|
sqlite_db = "/app/data/bodytrack.db"
|
|
profile_count = get_profile_count()
|
|
|
|
if os.path.exists(sqlite_db) and profile_count == 0:
|
|
print(" SQLite database found and PostgreSQL is empty")
|
|
print(" Starting automatic migration...")
|
|
# Import and run migration
|
|
try:
|
|
from migrate_to_postgres import main as migrate
|
|
migrate()
|
|
except Exception as e:
|
|
print(f"✗ Migration failed: {e}")
|
|
sys.exit(1)
|
|
elif os.path.exists(sqlite_db) and profile_count > 0:
|
|
print(f"⚠ SQLite DB exists but PostgreSQL already has {profile_count} profiles")
|
|
print(" Skipping migration (already migrated)")
|
|
elif not os.path.exists(sqlite_db):
|
|
print("✓ No SQLite database found (fresh install or already migrated)")
|
|
else:
|
|
print("✓ No migration needed")
|
|
|
|
print("\n✓ Database initialization complete")
|