"""Phase 3: billing fields, tenant_llm_keys, stripe_events, audit index, agent budget Revision ID: 005 Revises: 004 Create Date: 2026-03-24 This migration adds: 1. Billing columns on tenants table: - stripe_customer_id, stripe_subscription_id, stripe_subscription_item_id - subscription_status (TEXT, default 'none') - trial_ends_at (TIMESTAMPTZ, nullable) - agent_quota (INTEGER, default 0) 2. Budget column on agents table: - budget_limit_usd (FLOAT, nullable) — monthly spend cap per agent 3. tenant_llm_keys table: - Stores encrypted BYO API keys per tenant per provider - RLS enabled (same FORCE ROW LEVEL SECURITY pattern as agents) - UNIQUE(tenant_id, provider) constraint - key_hint column (VARCHAR(4)) for safe portal display without decryption - konstruct_app granted SELECT, INSERT, DELETE (no UPDATE — keys are immutable) 4. stripe_events table: - Idempotency guard for Stripe webhook event processing - Simple TEXT primary key (stripe event_id) - konstruct_app granted SELECT, INSERT (no UPDATE, no DELETE) 5. Composite index on audit_events: - idx_audit_events_tenant_type_created ON audit_events(tenant_id, action_type, created_at DESC) - Supports cost aggregation queries in usage.py endpoints Design: - tenant_llm_keys intentionally grants DELETE (portal operators can remove keys) - audit_events immutability is NOT weakened — only a new covering index is added - stripe_events does not need RLS — it's platform-wide idempotency, not tenant-scoped """ from __future__ import annotations from typing import Sequence, Union import sqlalchemy as sa from alembic import op from sqlalchemy.dialects.postgresql import UUID # revision identifiers, used by Alembic. revision: str = "005" down_revision: Union[str, None] = "004" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: # ========================================================================= # 1. Billing columns on tenants # ========================================================================= op.add_column("tenants", sa.Column( "stripe_customer_id", sa.String(255), nullable=True, comment="Stripe Customer ID (cus_...)", )) op.add_column("tenants", sa.Column( "stripe_subscription_id", sa.String(255), nullable=True, comment="Stripe Subscription ID (sub_...)", )) op.add_column("tenants", sa.Column( "stripe_subscription_item_id", sa.String(255), nullable=True, comment="Stripe Subscription Item ID (si_...) for quantity updates", )) op.add_column("tenants", sa.Column( "subscription_status", sa.String(50), nullable=False, server_default="none", comment="none | trialing | active | past_due | canceled | unpaid", )) op.add_column("tenants", sa.Column( "trial_ends_at", sa.DateTime(timezone=True), nullable=True, comment="Trial expiry timestamp (NULL for non-trial subscriptions)", )) op.add_column("tenants", sa.Column( "agent_quota", sa.Integer, nullable=False, server_default="0", comment="Number of active agents allowed under current subscription", )) # ========================================================================= # 2. Budget column on agents # ========================================================================= op.add_column("agents", sa.Column( "budget_limit_usd", sa.Float, nullable=True, comment="Monthly spend cap in USD. NULL means no limit.", )) # ========================================================================= # 3. tenant_llm_keys — encrypted BYO API keys # ========================================================================= op.create_table( "tenant_llm_keys", sa.Column( "id", UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()"), ), sa.Column( "tenant_id", UUID(as_uuid=True), sa.ForeignKey("tenants.id", ondelete="CASCADE"), nullable=False, ), sa.Column( "provider", sa.Text, nullable=False, comment="LLM provider: openai | anthropic | cohere | groq | etc.", ), sa.Column( "label", sa.Text, nullable=False, comment="Human-readable label for the portal", ), sa.Column( "encrypted_key", sa.Text, nullable=False, comment="Fernet-encrypted API key ciphertext", ), sa.Column( "key_hint", sa.String(4), nullable=True, comment="Last 4 chars of plaintext key for portal display", ), sa.Column( "key_version", sa.Integer, nullable=False, server_default="1", ), sa.Column( "created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("NOW()"), ), sa.UniqueConstraint("tenant_id", "provider", name="uq_tenant_llm_key_provider"), ) op.create_index("ix_tenant_llm_keys_tenant", "tenant_llm_keys", ["tenant_id"]) # RLS: only rows matching current tenant are visible op.execute("ALTER TABLE tenant_llm_keys ENABLE ROW LEVEL SECURITY") op.execute("ALTER TABLE tenant_llm_keys FORCE ROW LEVEL SECURITY") op.execute(""" CREATE POLICY tenant_isolation ON tenant_llm_keys USING (tenant_id = current_setting('app.current_tenant', TRUE)::uuid) """) # SELECT, INSERT, DELETE — no UPDATE (keys are immutable; rotate by delete+insert) op.execute("GRANT SELECT, INSERT, DELETE ON tenant_llm_keys TO konstruct_app") # ========================================================================= # 4. stripe_events — webhook idempotency guard # ========================================================================= op.create_table( "stripe_events", sa.Column( "event_id", sa.Text, primary_key=True, comment="Stripe event ID — globally unique per Stripe account", ), sa.Column( "processed_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("NOW()"), ), ) # No RLS — platform-wide idempotency table, not tenant-scoped op.execute("GRANT SELECT, INSERT ON stripe_events TO konstruct_app") # ========================================================================= # 5. Composite index on audit_events for usage aggregation queries # ========================================================================= # Covers: WHERE tenant_id = X AND action_type = 'llm_call' AND created_at >= ... # Used by: usage.py endpoints for per-agent and per-provider cost aggregation op.create_index( "idx_audit_events_tenant_type_created", "audit_events", ["tenant_id", "action_type", "created_at"], postgresql_ops={"created_at": "DESC"}, ) def downgrade() -> None: # Remove composite index on audit_events op.drop_index("idx_audit_events_tenant_type_created", table_name="audit_events") # Remove stripe_events op.execute("REVOKE ALL ON stripe_events FROM konstruct_app") op.drop_table("stripe_events") # Remove tenant_llm_keys op.execute("REVOKE ALL ON tenant_llm_keys FROM konstruct_app") op.drop_table("tenant_llm_keys") # Remove budget column from agents op.drop_column("agents", "budget_limit_usd") # Remove billing columns from tenants op.drop_column("tenants", "agent_quota") op.drop_column("tenants", "trial_ends_at") op.drop_column("tenants", "subscription_status") op.drop_column("tenants", "stripe_subscription_item_id") op.drop_column("tenants", "stripe_subscription_id") op.drop_column("tenants", "stripe_customer_id")