"""Billing production schema - auth onboarding, plans, packages, ledger, notifications. Revision ID: 003 Revises: 002 Create Date: 2025-03-10 Consolidates: auth onboarding + subscription_plans, billing tables, seeds, integrity constraints, billing_notifications. Server DB has 002 as last applied; all post-002 changes merged here. """ from typing import Sequence, Union import sqlalchemy as sa from alembic import op from sqlalchemy.dialects import postgresql revision: str = "003" down_revision: Union[str, None] = "002" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: # ---- subscription_plans (admin-managed plans with limits) ---- op.create_table( "subscription_plans", sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")), sa.Column("name", sa.String(100), nullable=False), sa.Column("slug", sa.String(50), unique=True, nullable=False), sa.Column("description", sa.Text()), sa.Column("is_active", sa.Boolean(), server_default="true"), sa.Column("price_monthly", sa.Integer(), server_default="0"), sa.Column("price_yearly", sa.Integer()), sa.Column("currency", sa.String(10), server_default="usd"), sa.Column("stripe_product_id", sa.String(100)), sa.Column("stripe_monthly_price_id", sa.String(100)), sa.Column("stripe_yearly_price_id", sa.String(100)), sa.Column("features", postgresql.JSONB()), sa.Column("limits", postgresql.JSONB()), sa.Column("sort_order", sa.Integer(), server_default="0"), sa.Column("trial_days", sa.Integer(), server_default="0"), sa.Column("includes_metered_billing", sa.Boolean(), server_default="false"), sa.Column("metered_stripe_price_id", sa.String(100)), sa.Column("metered_unit_label", sa.String(50)), sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), ) # ---- users: profile + last_login ---- op.add_column("users", sa.Column("first_name", sa.String(100))) op.add_column("users", sa.Column("last_name", sa.String(100))) op.add_column("users", sa.Column("phone", sa.String(50))) op.add_column("users", sa.Column("company", sa.String(255))) op.add_column("users", sa.Column("last_login_at", sa.TIMESTAMP(timezone=True))) # ---- workspaces: onboarding + security policy ---- op.add_column("workspaces", sa.Column("selected_plan_slug", sa.String(50))) op.add_column("workspaces", sa.Column("onboarding_status", sa.String(50))) op.add_column( "workspaces", sa.Column("security_policy_require_mfa_for_paid_users", sa.Boolean(), server_default="true"), ) # ---- subscriptions: plan_id, plan_name (legacy), stripe identity ---- op.add_column("subscriptions", sa.Column("plan_id", sa.String(100))) op.add_column("subscriptions", sa.Column("plan_name", sa.String(100))) op.add_column("subscriptions", sa.Column("stripe_customer_id", sa.String(100))) op.add_column("subscriptions", sa.Column("trial_end", sa.TIMESTAMP(timezone=True))) op.add_column("subscriptions", sa.Column("metered_subscription_item_id", sa.String(100))) # ---- workspace_invites: status lifecycle ---- op.add_column("workspace_invites", sa.Column("status", sa.String(20), server_default="pending")) op.add_column("workspace_invites", sa.Column("revoked_at", sa.TIMESTAMP(timezone=True))) # ---- Seed free plan ---- op.execute(""" INSERT INTO subscription_plans ( id, name, slug, description, is_active, price_monthly, price_yearly, currency, limits, sort_order, trial_days ) SELECT gen_random_uuid(), 'Free', 'free', 'Free tier for getting started', true, 0, 0, 'usd', '{"max_members": 1}'::jsonb, 0, 0 WHERE NOT EXISTS (SELECT 1 FROM subscription_plans WHERE slug = 'free') """) # ---- subscriptions: plan_slug, stripe_product_id (stripe_price_id exists from 001) ---- op.add_column("subscriptions", sa.Column("plan_slug", sa.String(50))) op.add_column("subscriptions", sa.Column("stripe_product_id", sa.String(100))) op.add_column("subscriptions", sa.Column("billing_interval", sa.String(20))) op.add_column("subscriptions", sa.Column("cancel_at_period_end", sa.Boolean(), server_default="false")) op.add_column("subscriptions", sa.Column("trial_start", sa.TIMESTAMP(timezone=True))) op.add_column("subscriptions", sa.Column("ended_at", sa.TIMESTAMP(timezone=True))) op.add_column("subscriptions", sa.Column("metadata_json", postgresql.JSONB())) op.add_column( "subscriptions", sa.Column("subscription_plan_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("subscription_plans.id")), ) op.execute(""" UPDATE subscriptions SET stripe_price_id = plan_id WHERE plan_id IS NOT NULL AND (stripe_price_id IS NULL OR stripe_price_id = '') """) op.execute(""" UPDATE subscriptions SET billing_interval = 'month' WHERE billing_interval IS NULL """) # ---- usage_credit_packages ---- op.create_table( "usage_credit_packages", sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")), sa.Column("slug", sa.String(50), unique=True, nullable=False), sa.Column("name", sa.String(100), nullable=False), sa.Column("description", sa.Text()), sa.Column("units", sa.Integer(), nullable=False), sa.Column("price_minor", sa.BigInteger(), nullable=False), sa.Column("currency", sa.String(10), server_default="usd"), sa.Column("stripe_product_id", sa.String(100)), sa.Column("stripe_price_id", sa.String(100)), sa.Column("is_active", sa.Boolean(), server_default="true"), sa.Column("sort_order", sa.Integer(), server_default="0"), sa.Column("metadata_json", postgresql.JSONB()), sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), ) # ---- usage_credit_purchases ---- op.create_table( "usage_credit_purchases", sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")), sa.Column("workspace_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("workspaces.id", ondelete="CASCADE"), nullable=False), sa.Column("stripe_customer_id", sa.String(100), nullable=False), sa.Column("stripe_checkout_session_id", sa.String(100)), sa.Column("stripe_payment_intent_id", sa.String(100)), sa.Column("stripe_invoice_id", sa.String(100)), sa.Column("usage_credit_package_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("usage_credit_packages.id")), sa.Column("package_slug", sa.String(50), nullable=False), sa.Column("units_purchased", sa.Integer(), nullable=False), sa.Column("amount_minor", sa.BigInteger(), nullable=False), sa.Column("currency", sa.String(10), server_default="usd"), sa.Column("status", sa.String(30), nullable=False), sa.Column("paid_at", sa.TIMESTAMP(timezone=True)), sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), ) op.create_index("idx_usage_credit_purchases_workspace", "usage_credit_purchases", ["workspace_id"]) op.create_index("idx_usage_credit_purchases_status", "usage_credit_purchases", ["status"]) # ---- workspace_credit_ledger ---- op.create_table( "workspace_credit_ledger", sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")), sa.Column("workspace_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("workspaces.id", ondelete="CASCADE"), nullable=False), sa.Column("entry_type", sa.String(30), nullable=False), sa.Column("units_delta", sa.Integer(), nullable=False), sa.Column("balance_after", sa.Integer()), sa.Column("source_bucket", sa.String(50)), sa.Column("source_type", sa.String(50)), sa.Column("source_id", sa.String(100)), sa.Column("subscription_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("subscriptions.id")), sa.Column("usage_credit_purchase_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("usage_credit_purchases.id")), sa.Column("project_id", postgresql.UUID(as_uuid=True)), sa.Column("codegen_job_id", postgresql.UUID(as_uuid=True)), sa.Column("artifact_id", postgresql.UUID(as_uuid=True)), sa.Column("revision_id", postgresql.UUID(as_uuid=True)), sa.Column("description", sa.Text()), sa.Column("metadata_json", postgresql.JSONB()), sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), ) op.create_index("idx_workspace_credit_ledger_workspace", "workspace_credit_ledger", ["workspace_id", "created_at"]) # ---- workspace_usage_allowances ---- op.create_table( "workspace_usage_allowances", sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")), sa.Column("workspace_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("workspaces.id", ondelete="CASCADE"), nullable=False), sa.Column("subscription_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("subscriptions.id")), sa.Column("period_start", sa.TIMESTAMP(timezone=True), nullable=False), sa.Column("period_end", sa.TIMESTAMP(timezone=True), nullable=False), sa.Column("included_units", sa.Integer(), nullable=False), sa.Column("consumed_units", sa.Integer(), server_default="0"), sa.Column("remaining_units", sa.Integer(), nullable=False), sa.Column("reset_strategy", sa.String(30), server_default="period_end"), sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), ) op.create_index("idx_workspace_usage_allowances_workspace", "workspace_usage_allowances", ["workspace_id", "period_end"]) # ---- Seed Basic, Pro, Enterprise ---- op.execute(""" INSERT INTO subscription_plans (id, name, slug, description, is_active, price_monthly, price_yearly, currency, limits, features, sort_order, trial_days) SELECT gen_random_uuid(), 'Basic', 'basic', 'For individuals - 3 projects, 1 member', true, 2000, NULL, 'usd', '{"max_projects": 3, "max_members": 1, "included_ai_units": 100}'::jsonb, '{"custom_domain": false, "extra_staff": false}'::jsonb, 1, 0 WHERE NOT EXISTS (SELECT 1 FROM subscription_plans WHERE slug = 'basic') """) op.execute(""" INSERT INTO subscription_plans (id, name, slug, description, is_active, price_monthly, price_yearly, currency, limits, features, sort_order, trial_days) SELECT gen_random_uuid(), 'Pro', 'pro', 'For teams - 10 projects, custom domain, 3x AI allowance', true, 6000, 60000, 'usd', '{"max_projects": 10, "max_members": 10, "included_ai_units": 500, "ai_multiplier": 3}'::jsonb, '{"custom_domain": true, "extra_staff": true}'::jsonb, 2, 0 WHERE NOT EXISTS (SELECT 1 FROM subscription_plans WHERE slug = 'pro') """) op.execute(""" INSERT INTO subscription_plans (id, name, slug, description, is_active, price_monthly, price_yearly, currency, limits, features, sort_order, trial_days) SELECT gen_random_uuid(), 'Enterprise', 'enterprise', 'Unlimited or higher caps, enterprise features', true, 20000, NULL, 'usd', '{"max_projects": 999, "max_members": 999, "included_ai_units": 5000}'::jsonb, '{"custom_domain": true, "extra_staff": true, "enterprise": true}'::jsonb, 3, 0 WHERE NOT EXISTS (SELECT 1 FROM subscription_plans WHERE slug = 'enterprise') """) # ---- Seed usage packages ---- op.execute(""" INSERT INTO usage_credit_packages (id, slug, name, description, units, price_minor, currency, is_active, sort_order) SELECT gen_random_uuid(), 'minutes_60', '60 AI Minutes', '60 additional AI generation minutes', 60, 999, 'usd', true, 1 WHERE NOT EXISTS (SELECT 1 FROM usage_credit_packages WHERE slug = 'minutes_60') """) op.execute(""" INSERT INTO usage_credit_packages (id, slug, name, description, units, price_minor, currency, is_active, sort_order) SELECT gen_random_uuid(), 'minutes_180', '180 AI Minutes', '180 additional AI generation minutes', 180, 2499, 'usd', true, 2 WHERE NOT EXISTS (SELECT 1 FROM usage_credit_packages WHERE slug = 'minutes_180') """) op.execute(""" INSERT INTO usage_credit_packages (id, slug, name, description, units, price_minor, currency, is_active, sort_order) SELECT gen_random_uuid(), 'minutes_500', '500 AI Minutes', '500 additional AI generation minutes', 500, 5999, 'usd', true, 3 WHERE NOT EXISTS (SELECT 1 FROM usage_credit_packages WHERE slug = 'minutes_500') """) # ---- Billing integrity: unique checkout session, indexes ---- op.execute(""" CREATE UNIQUE INDEX uq_usage_credit_purchases_stripe_session ON usage_credit_purchases (stripe_checkout_session_id) WHERE stripe_checkout_session_id IS NOT NULL """) op.create_index("idx_usage_credit_packages_stripe_price_id", "usage_credit_packages", ["stripe_price_id"], unique=False) op.create_index("idx_subscriptions_stripe_price_id", "subscriptions", ["stripe_price_id"], unique=False) op.create_index("idx_subscriptions_plan_slug", "subscriptions", ["plan_slug"], unique=False) op.create_index("uq_workspace_usage_allowances_workspace_period", "workspace_usage_allowances", ["workspace_id", "period_start", "period_end"], unique=True) # ---- billing_notifications ---- op.create_table( "billing_notifications", sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")), sa.Column("workspace_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("workspaces.id", ondelete="CASCADE"), nullable=False), sa.Column("notification_type", sa.String(50), nullable=False), sa.Column("source_object_type", sa.String(50), nullable=False), sa.Column("source_object_id", sa.String(100), nullable=False), sa.Column("dedupe_key", sa.String(200), nullable=False), sa.Column("recipient_email", sa.String(255), nullable=False), sa.Column("status", sa.String(30), nullable=False), sa.Column("provider_message_id", sa.String(200)), sa.Column("last_error", sa.Text()), sa.Column("created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), sa.Column("sent_at", sa.TIMESTAMP(timezone=True)), sa.Column("updated_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now()), ) op.create_index("idx_billing_notifications_workspace", "billing_notifications", ["workspace_id", "created_at"]) op.create_index("idx_billing_notifications_dedupe", "billing_notifications", ["dedupe_key", "recipient_email"], unique=True) def downgrade() -> None: op.drop_index("idx_billing_notifications_dedupe", table_name="billing_notifications") op.drop_index("idx_billing_notifications_workspace", table_name="billing_notifications") op.drop_table("billing_notifications") op.drop_index("uq_workspace_usage_allowances_workspace_period", table_name="workspace_usage_allowances") op.drop_index("idx_subscriptions_plan_slug", table_name="subscriptions") op.drop_index("idx_subscriptions_stripe_price_id", table_name="subscriptions") op.drop_index("idx_usage_credit_packages_stripe_price_id", table_name="usage_credit_packages") op.execute("DROP INDEX IF EXISTS uq_usage_credit_purchases_stripe_session") op.execute("DELETE FROM usage_credit_packages WHERE slug IN ('minutes_60', 'minutes_180', 'minutes_500')") op.execute("DELETE FROM subscription_plans WHERE slug IN ('basic', 'pro', 'enterprise')") op.drop_index("idx_workspace_usage_allowances_workspace", table_name="workspace_usage_allowances") op.drop_table("workspace_usage_allowances") op.drop_index("idx_workspace_credit_ledger_workspace", table_name="workspace_credit_ledger") op.drop_table("workspace_credit_ledger") op.drop_index("idx_usage_credit_purchases_status", table_name="usage_credit_purchases") op.drop_index("idx_usage_credit_purchases_workspace", table_name="usage_credit_purchases") op.drop_table("usage_credit_purchases") op.drop_table("usage_credit_packages") op.drop_column("subscriptions", "subscription_plan_id") op.drop_column("subscriptions", "metadata_json") op.drop_column("subscriptions", "ended_at") op.drop_column("subscriptions", "trial_start") op.drop_column("subscriptions", "cancel_at_period_end") op.drop_column("subscriptions", "billing_interval") op.drop_column("subscriptions", "stripe_product_id") op.drop_column("subscriptions", "plan_slug") op.execute("DELETE FROM subscription_plans WHERE slug = 'free'") op.drop_column("workspace_invites", "revoked_at") op.drop_column("workspace_invites", "status") op.drop_column("subscriptions", "metered_subscription_item_id") op.drop_column("subscriptions", "trial_end") op.drop_column("subscriptions", "stripe_customer_id") op.drop_column("subscriptions", "plan_name") op.drop_column("subscriptions", "plan_id") op.drop_column("workspaces", "security_policy_require_mfa_for_paid_users") op.drop_column("workspaces", "onboarding_status") op.drop_column("workspaces", "selected_plan_slug") op.drop_column("users", "last_login_at") op.drop_column("users", "company") op.drop_column("users", "phone") op.drop_column("users", "last_name") op.drop_column("users", "first_name") op.drop_table("subscription_plans")