feat(03-01): DB migrations, models, encryption service, and test scaffolds
- 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)
This commit is contained in:
234
migrations/versions/005_billing_and_usage.py
Normal file
234
migrations/versions/005_billing_and_usage.py
Normal file
@@ -0,0 +1,234 @@
|
||||
"""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")
|
||||
Reference in New Issue
Block a user