init_db.py

import os
import sys
import psycopg2
from dotenv import load_dotenv

def main():
    """Initialize the PostgreSQL database"""
    # Load environment variables
    load_dotenv()
    
    # Get database connection info from environment variables
    database_url = os.getenv('DATABASE_URL')
    if not database_url:
        print("No DATABASE_URL found in environment. Using default connection info.")
        db_user = 'postgres'
        db_pass = 'postgres'
        db_host = 'localhost'
        db_port = '5432'
    else:
        # Parse DATABASE_URL
        # Format: postgresql://username:password@hostname:port/database
        try:
            # Remove postgresql:// prefix
            db_info = database_url.split('://', 1)[1]
            # Split user:pass and host:port/dbname
            auth, db_path = db_info.split('@', 1)
            # Get username and password
            db_user, db_pass = auth.split(':', 1)
            # Get host, port, and dbname
            db_host_port, db_name = db_path.split('/', 1)
            if ':' in db_host_port:
                db_host, db_port = db_host_port.split(':', 1)
            else:
                db_host = db_host_port
                db_port = '5432'
        except Exception as e:
            print(f"Error parsing DATABASE_URL: {e}")
            return 1
    
    # Database name
    db_name = 'sas_db'
    
    # Connect to PostgreSQL server
    try:
        conn = psycopg2.connect(
            host=db_host,
            port=db_port,
            user=db_user,
            password=db_pass,
            dbname='postgres'  # Connect to default DB first
        )
        conn.autocommit = True
        cursor = conn.cursor()
        
        # Check if database exists
        cursor.execute("SELECT 1 FROM pg_catalog.pg_database WHERE datname = %s", (db_name,))
        exists = cursor.fetchone()
        
        if not exists:
            print(f"Creating database '{db_name}'...")
            cursor.execute(f"CREATE DATABASE {db_name}")
            print(f"Database '{db_name}' created successfully.")
        else:
            print(f"Database '{db_name}' already exists.")
        
        # Close connection to postgres
        cursor.close()
        conn.close()
        
        print("Database initialized successfully.")
        print("Run 'python app.py' to start the application.")
        return 0
        
    except Exception as e:
        print(f"Error initializing database: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())