"""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")