resolutionflow/backend/alembic/versions/a1b2c3d4e5f6_add_ai_flow_builder.py

"""add ai flow builder tables and columns

Revision ID: a1b2c3d4e5f6
Revises: e65b9f8fd458
Create Date: 2026-02-20 12:00:00.000000

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql


revision: str = "a1b2c3d4e5f6"
down_revision: Union[str, None] = "e65b9f8fd458"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ── ai_conversations table ──
    op.create_table(
        "ai_conversations",
        sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column(
            "user_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("users.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "account_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("accounts.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("status", sa.String(20), nullable=False, server_default="foundation"),
        sa.Column("messages", postgresql.JSONB(), nullable=False, server_default="[]"),
        sa.Column(
            "wizard_state", postgresql.JSONB(), nullable=False, server_default="{}"
        ),
        sa.Column("generated_tree", postgresql.JSONB(), nullable=True),
        sa.Column("question_rounds", sa.Integer(), nullable=False, server_default="0"),
        sa.Column(
            "expires_at", sa.DateTime(timezone=True), nullable=False
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
    )
    op.create_index(
        "ix_ai_conversations_user_id", "ai_conversations", ["user_id"]
    )
    op.create_index(
        "ix_ai_conversations_account_id", "ai_conversations", ["account_id"]
    )
    op.create_index(
        "ix_ai_conversations_user_created",
        "ai_conversations",
        ["user_id", sa.text("created_at DESC")],
    )
    op.create_index(
        "ix_ai_conversations_expires_at", "ai_conversations", ["expires_at"]
    )

    # ── ai_usage table ──
    op.create_table(
        "ai_usage",
        sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column(
            "user_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("users.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "account_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("accounts.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "conversation_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("ai_conversations.id", ondelete="SET NULL"),
            nullable=True,
        ),
        sa.Column("generation_type", sa.String(20), nullable=False),
        sa.Column("tier_at_time", sa.String(20), nullable=False),
        sa.Column("input_tokens", sa.Integer(), nullable=False, server_default="0"),
        sa.Column("output_tokens", sa.Integer(), nullable=False, server_default="0"),
        sa.Column(
            "estimated_cost_usd",
            sa.Numeric(10, 6),
            nullable=False,
            server_default="0",
        ),
        sa.Column("succeeded", sa.Boolean(), nullable=False, server_default="true"),
        sa.Column(
            "counts_toward_quota",
            sa.Boolean(),
            nullable=False,
            server_default="false",
        ),
        sa.Column("error_code", sa.String(100), nullable=True),
        sa.Column("metadata", postgresql.JSONB(), nullable=False, server_default="{}"),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
    )
    op.create_index("ix_ai_usage_user_id", "ai_usage", ["user_id"])
    op.create_index("ix_ai_usage_account_id", "ai_usage", ["account_id"])
    op.create_index("ix_ai_usage_created_at", "ai_usage", ["created_at"])
    op.create_index(
        "ix_ai_usage_user_created",
        "ai_usage",
        ["user_id", sa.text("created_at DESC")],
    )
    op.create_index(
        "ix_ai_usage_user_type_created",
        "ai_usage",
        ["user_id", "generation_type", sa.text("created_at DESC")],
    )
    # Prevents double quota decrement from race conditions
    op.execute(
        """
        CREATE UNIQUE INDEX ix_ai_usage_unique_quota
        ON ai_usage (conversation_id)
        WHERE counts_toward_quota = true;
        """
    )

    # ── Schema modifications to existing tables ──

    # users: add ai_billing_cycle_anchor_at
    op.add_column(
        "users",
        sa.Column("ai_billing_cycle_anchor_at", sa.DateTime(timezone=True), nullable=True),
    )
    # Backfill: use created_at as the billing anchor
    op.execute(
        "UPDATE users SET ai_billing_cycle_anchor_at = created_at WHERE ai_billing_cycle_anchor_at IS NULL"
    )

    # plan_limits: add AI limit columns
    op.add_column(
        "plan_limits",
        sa.Column("max_ai_builds_per_month", sa.Integer(), nullable=True),
    )
    op.add_column(
        "plan_limits",
        sa.Column("max_ai_builds_per_24h", sa.Integer(), nullable=True),
    )

    # account_limit_overrides: add AI override columns
    op.add_column(
        "account_limit_overrides",
        sa.Column("override_max_ai_builds_per_month", sa.Integer(), nullable=True),
    )
    op.add_column(
        "account_limit_overrides",
        sa.Column("override_max_ai_builds_per_24h", sa.Integer(), nullable=True),
    )

    # Seed plan_limits with AI quota values
    op.execute(
        """
        UPDATE plan_limits SET max_ai_builds_per_month = 2, max_ai_builds_per_24h = 1
        WHERE plan = 'free';
        """
    )
    op.execute(
        """
        UPDATE plan_limits SET max_ai_builds_per_month = 50, max_ai_builds_per_24h = 10
        WHERE plan = 'pro';
        """
    )
    op.execute(
        """
        UPDATE plan_limits SET max_ai_builds_per_month = 200, max_ai_builds_per_24h = 20
        WHERE plan = 'team';
        """
    )
    # Enterprise: NULL means unlimited (no update needed as default is NULL)


def downgrade() -> None:
    # Drop AI override columns from account_limit_overrides
    op.drop_column("account_limit_overrides", "override_max_ai_builds_per_24h")
    op.drop_column("account_limit_overrides", "override_max_ai_builds_per_month")

    # Drop AI limit columns from plan_limits
    op.drop_column("plan_limits", "max_ai_builds_per_24h")
    op.drop_column("plan_limits", "max_ai_builds_per_month")

    # Drop ai_billing_cycle_anchor_at from users
    op.drop_column("users", "ai_billing_cycle_anchor_at")

    # Drop ai_usage table (indexes drop automatically)
    op.drop_table("ai_usage")

    # Drop ai_conversations table (indexes drop automatically)
    op.drop_table("ai_conversations")