data-entry-app/backend/app/db/migrations.py

from __future__ import annotations

from sqlalchemy import inspect, text
from sqlalchemy.engine import Engine


TENANT_TABLES = {
    "client_users": None,
    "client_feature_access": None,
    "raw_materials": None,
    "raw_material_price_versions": None,
    "mixes": None,
    "mix_ingredients": None,
    "products": None,
    "scenarios": None,
    "costing_results": None,
    "process_cost_rules": None,
    "packaging_cost_rules": None,
    "freight_cost_rules": None,
}


def _has_column(engine: Engine, table_name: str, column_name: str) -> bool:
    inspector = inspect(engine)
    try:
        columns = inspector.get_columns(table_name)
    except Exception:
        return False
    return any(column["name"] == column_name for column in columns)


def _add_tenant_column(engine: Engine, table_name: str) -> None:
    if _has_column(engine, table_name, "tenant_id"):
        return
    with engine.begin() as connection:
        connection.execute(text(f"ALTER TABLE {table_name} ADD COLUMN tenant_id VARCHAR(64)"))


def _table_exists(engine: Engine, table_name: str) -> bool:
    return inspect(engine).has_table(table_name)


def ensure_tenant_columns(engine: Engine) -> None:
    for table_name in TENANT_TABLES:
        if _table_exists(engine, table_name):
            _add_tenant_column(engine, table_name)


def sync_tenant_ids(engine: Engine) -> None:
    if not _table_exists(engine, "client_accounts"):
        return

    with engine.begin() as connection:
        default_tenant = connection.execute(
            text("SELECT tenant_id FROM client_accounts ORDER BY id LIMIT 1")
        ).scalar_one_or_none()
        if not default_tenant:
            return

        statements = [
            text(
                """
                UPDATE client_users
                SET tenant_id = (
                    SELECT client_accounts.tenant_id
                    FROM client_accounts
                    WHERE client_accounts.id = client_users.client_account_id
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE client_feature_access
                SET tenant_id = (
                    SELECT client_accounts.tenant_id
                    FROM client_accounts
                    WHERE client_accounts.id = client_feature_access.client_account_id
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE raw_materials
                SET tenant_id = :default_tenant
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE raw_material_price_versions
                SET tenant_id = (
                    SELECT raw_materials.tenant_id
                    FROM raw_materials
                    WHERE raw_materials.id = raw_material_price_versions.raw_material_id
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE mixes
                SET tenant_id = COALESCE(
                    (
                        SELECT client_accounts.tenant_id
                        FROM client_accounts
                        WHERE client_accounts.name = mixes.client_name
                    ),
                    :default_tenant
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE mix_ingredients
                SET tenant_id = (
                    SELECT mixes.tenant_id
                    FROM mixes
                    WHERE mixes.id = mix_ingredients.mix_id
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE products
                SET tenant_id = COALESCE(
                    (
                        SELECT client_accounts.tenant_id
                        FROM client_accounts
                        WHERE client_accounts.name = products.client_name
                    ),
                    (
                        SELECT mixes.tenant_id
                        FROM mixes
                        WHERE mixes.id = products.mix_id
                    ),
                    :default_tenant
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE scenarios
                SET tenant_id = :default_tenant
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE costing_results
                SET tenant_id = COALESCE(
                    (
                        SELECT products.tenant_id
                        FROM products
                        WHERE products.id = costing_results.product_id
                    ),
                    (
                        SELECT scenarios.tenant_id
                        FROM scenarios
                        WHERE scenarios.id = costing_results.scenario_id
                    ),
                    :default_tenant
                )
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE process_cost_rules
                SET tenant_id = :default_tenant
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE packaging_cost_rules
                SET tenant_id = :default_tenant
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
            text(
                """
                UPDATE freight_cost_rules
                SET tenant_id = :default_tenant
                WHERE tenant_id IS NULL OR tenant_id = '' OR tenant_id = 'default'
                """
            ),
        ]

        for statement in statements:
            connection.execute(statement, {"default_tenant": default_tenant})