mitai-jinkendo/backend/db_init.py

#!/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 (v9b)")
    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")