Adolfo Delorenzo
215e67a7eb
- Add stripe and cryptography to shared pyproject.toml - Add recharts, @stripe/stripe-js, stripe to portal package.json (submodule) - Add billing fields to Tenant model (stripe_customer_id, subscription_status, agent_quota, trial_ends_at) - Add budget_limit_usd to Agent model - Create TenantLlmKey and StripeEvent models in billing.py (AuditBase and Base respectively) - Create KeyEncryptionService (MultiFernet encrypt/decrypt/rotate) in crypto.py - Create compute_budget_status helper in usage.py (threshold logic: ok/warning/exceeded) - Add platform_encryption_key, stripe_, slack_oauth settings to config.py - Create Alembic migration 005 with all schema changes, RLS, grants, and composite index - All 12 tests passing (key encryption roundtrip, rotation, budget thresholds)
235 lines
8.1 KiB
Python
235 lines
8.1 KiB
Python
"""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")
|