migrate.py

import os
import bcrypt
from datetime import datetime
from sqlalchemy import text, inspect
from database import engine

def run_migrations():
    print("Running database migrations...")
    run_engine_migrations(engine)

def run_db_migrations(engine_param):
    print("Running SQLAlchemy engine migrations...")
    run_engine_migrations(engine_param)

def run_engine_migrations(engine):
    # Get database inspector
    inspector = inspect(engine)
    
    # 0. Seed default user if not exists (in users table)
    if "users" in inspector.get_table_names():
        try:
            with engine.connect() as conn:
                result = conn.execute(text("SELECT id FROM users WHERE id = 1")).fetchone()
                if not result:
                    hashed_pass = bcrypt.hashpw(b"admin123", bcrypt.gensalt()).decode('utf-8')
                    conn.execute(
                        text("INSERT INTO users (id, email, hashed_password, company_name, created_at) VALUES (1, :email, :password, :company, :created)"),
                        {
                            "email": "admin@founder.os",
                            "password": hashed_pass,
                            "company": "Erha Technologies",
                            "created": datetime.utcnow()
                        }
                    )
                    conn.commit()
                    print("Migration: Inserted default admin user.")
        except Exception as e:
            print(f"Error seeding default user: {e}")

    # 1. Migrate clients table
    if "clients" in inspector.get_table_names():
        columns_clients = [c['name'] for c in inspector.get_columns('clients')]
        migrations_clients = [
            ("description", "TEXT"),
            ("is_active", "BOOLEAN DEFAULT TRUE"),
            ("advance_payment", "FLOAT DEFAULT 0.0"),
            ("remaining_payment", "FLOAT DEFAULT 0.0"),
            ("joining_date", "TIMESTAMP"),
            ("company", "VARCHAR(255)"),
            ("address", "TEXT")
        ]
        for col_name, col_type in migrations_clients:
            if col_name not in columns_clients:
                try:
                    with engine.begin() as conn:
                        conn.execute(text(f"ALTER TABLE clients ADD COLUMN {col_name} {col_type}"))
                    print(f"Migration: Added column '{col_name}' ({col_type}) to 'clients' table.")
                except Exception as e:
                    print(f"Error adding '{col_name}' to 'clients': {e}")

    # 2. Migrate employees table
    if "employees" in inspector.get_table_names():
        columns_employees = [c['name'] for c in inspector.get_columns('employees')]
        migrations_employees = [
            ("bank_account_number", "VARCHAR(255)"),
            ("bank_name", "VARCHAR(255)"),
            ("account_title", "VARCHAR(255)"),
            ("starting_salary", "FLOAT DEFAULT 0.0")
        ]
        for col_name, col_type in migrations_employees:
            if col_name not in columns_employees:
                try:
                    with engine.begin() as conn:
                        conn.execute(text(f"ALTER TABLE employees ADD COLUMN {col_name} {col_type}"))
                    print(f"Migration: Added column '{col_name}' ({col_type}) to 'employees' table.")
                    if col_name == "starting_salary":
                        with engine.begin() as conn:
                            conn.execute(text("UPDATE employees SET starting_salary = base_salary WHERE starting_salary IS NULL OR starting_salary = 0.0"))
                except Exception as e:
                    print(f"Error adding '{col_name}' to 'employees': {e}")

    # 3. Add user_id column to all entity tables for multi-tenancy
    tables_to_migrate = ["clients", "employees", "bank_accounts", "invoices", "payroll", "revenues", "expenses", "security_config", "personal_expenses"]
    for table in tables_to_migrate:
        if table in inspector.get_table_names():
            columns = [c['name'] for c in inspector.get_columns(table)]
            if "user_id" not in columns:
                try:
                    with engine.begin() as conn:
                        conn.execute(text(f"ALTER TABLE {table} ADD COLUMN user_id INTEGER DEFAULT 1"))
                    print(f"Migration: Added 'user_id' column to '{table}' table.")
                except Exception as e:
                    print(f"Error adding 'user_id' to '{table}': {e}")

    # 4. Add employee_id, payroll_id and created_at columns to expenses table
    if "expenses" in inspector.get_table_names():
        columns_expenses = [c['name'] for c in inspector.get_columns('expenses')]
        migrations_expenses = [
            ("employee_id", "INTEGER"),
            ("payroll_id", "INTEGER"),
            ("created_at", "TIMESTAMP"),
            ("target_month", "VARCHAR(50)"),
            ("target_year", "INTEGER")
        ]
        for col_name, col_type in migrations_expenses:
            if col_name not in columns_expenses:
                try:
                    with engine.begin() as conn:
                        conn.execute(text(f"ALTER TABLE expenses ADD COLUMN {col_name} {col_type}"))
                    print(f"Migration: Added column '{col_name}' ({col_type}) to 'expenses' table.")
                except Exception as e:
                    print(f"Error adding '{col_name}' to 'expenses': {e}")

    # 5. Add invoice_id, payroll_id, expense_id, revenue_id columns to transactions table
    if "transactions" in inspector.get_table_names():
        columns_txn = [c['name'] for c in inspector.get_columns('transactions')]
        migrations_tx = [
            ("invoice_id", "INTEGER"),
            ("payroll_id", "INTEGER"),
            ("expense_id", "INTEGER"),
            ("revenue_id", "INTEGER")
        ]
        for col_name, col_type in migrations_tx:
            if col_name not in columns_txn:
                try:
                    with engine.begin() as conn:
                        conn.execute(text(f"ALTER TABLE transactions ADD COLUMN {col_name} {col_type}"))
                    print(f"Migration: Added column '{col_name}' ({col_type}) to 'transactions' table.")
                except Exception as e:
                    print(f"Error adding '{col_name}' to 'transactions': {e}")

    # 6. Migrate salary_increments table
    if "salary_increments" in inspector.get_table_names():
        columns_inc = [c['name'] for c in inspector.get_columns('salary_increments')]
        migrations_inc = [
            ("old_salary", "FLOAT DEFAULT 0.0"),
            ("new_salary", "FLOAT DEFAULT 0.0"),
            ("increment_amount", "FLOAT DEFAULT 0.0"),
            ("date", "TIMESTAMP")
        ]
        for col_name, col_type in migrations_inc:
            if col_name not in columns_inc:
                try:
                    with engine.begin() as conn:
                        conn.execute(text(f"ALTER TABLE salary_increments ADD COLUMN {col_name} {col_type}"))
                    print(f"Migration: Added column '{col_name}' ({col_type}) to 'salary_increments' table.")
                except Exception as e:
                    print(f"Error adding '{col_name}' to 'salary_increments': {e}")

    print("Migrations completed successfully.")

if __name__ == "__main__":
    run_migrations()