"""initial gateway schema Creates schema ``gateway``, the three enum types, all tables and indexes, and the ``notify_key_revoked()`` function plus ``trg_notify_key_revoked`` trigger, matching SPEC §5 verbatim in structure. Revision ID: 0001_initial Revises: Create Date: 2026-05-22 """ from __future__ import annotations from collections.abc import Sequence import sqlalchemy as sa from alembic import op from sqlalchemy.dialects import postgresql # revision identifiers, used by Alembic. revision: str = "0001_initial" down_revision: str | None = None branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None SCHEMA = "gateway" # Enum types are created explicitly via raw SQL below; the table columns # reference them with create_type=False so they are not created twice. _key_status = postgresql.ENUM( "active", "disabled", "revoked", name="key_status", schema=SCHEMA, create_type=False ) _tenant_status = postgresql.ENUM( "active", "suspended", "closed", name="tenant_status", schema=SCHEMA, create_type=False ) _budget_period = postgresql.ENUM( "day", "month", "total", name="budget_period", schema=SCHEMA, create_type=False ) def upgrade() -> None: """Create the full ``gateway`` schema.""" op.execute(f'CREATE SCHEMA IF NOT EXISTS "{SCHEMA}"') # --- Enum types (SPEC §5) --- op.execute("CREATE TYPE gateway.key_status AS ENUM ('active', 'disabled', 'revoked')") op.execute("CREATE TYPE gateway.tenant_status AS ENUM ('active', 'suspended', 'closed')") op.execute("CREATE TYPE gateway.budget_period AS ENUM ('day', 'month', 'total')") # --- tenants --- op.create_table( "tenants", sa.Column( "id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()"), ), sa.Column("name", sa.Text(), nullable=False, unique=True), sa.Column( "status", _tenant_status, nullable=False, server_default=sa.text("'active'") ), sa.Column( "created_at", postgresql.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("now()"), ), sa.Column( "metadata", postgresql.JSONB(), nullable=False, server_default=sa.text("'{}'::jsonb"), ), schema=SCHEMA, ) # --- tenant_limits --- op.create_table( "tenant_limits", sa.Column( "tenant_id", postgresql.UUID(as_uuid=True), sa.ForeignKey(f"{SCHEMA}.tenants.id", ondelete="CASCADE"), primary_key=True, ), sa.Column("rpm", sa.Integer(), nullable=False, server_default=sa.text("60")), sa.Column("tpm", sa.Integer(), nullable=False, server_default=sa.text("100000")), sa.Column("concurrent", sa.Integer(), nullable=False, server_default=sa.text("8")), sa.Column("tokens_daily", sa.BigInteger(), nullable=True), sa.Column("tokens_monthly", sa.BigInteger(), nullable=True), sa.Column("tokens_total", sa.BigInteger(), nullable=True), sa.Column( "allowed_models", postgresql.ARRAY(sa.Text()), nullable=False, server_default=sa.text("'{}'"), ), sa.Column( "allow_all_models", sa.Boolean(), nullable=False, server_default=sa.text("false"), ), sa.Column( "log_prompts_default", sa.Boolean(), nullable=False, server_default=sa.text("false"), ), sa.Column( "prompt_retention_days", sa.Integer(), nullable=False, server_default=sa.text("30") ), sa.Column( "audit_retention_days", sa.Integer(), nullable=False, server_default=sa.text("365") ), schema=SCHEMA, ) # --- api_keys --- op.create_table( "api_keys", sa.Column( "id", postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()"), ), sa.Column( "tenant_id", postgresql.UUID(as_uuid=True), sa.ForeignKey(f"{SCHEMA}.tenants.id", ondelete="CASCADE"), nullable=False, ), sa.Column("prefix", sa.Text(), nullable=False, unique=True), sa.Column("key_hash", sa.Text(), nullable=False), sa.Column("name", sa.Text(), nullable=False), sa.Column("status", _key_status, nullable=False, server_default=sa.text("'active'")), sa.Column( "scopes", postgresql.ARRAY(sa.Text()), nullable=False, server_default=sa.text("'{chat,embeddings}'"), ), sa.Column( "created_at", postgresql.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("now()"), ), sa.Column("last_used_at", postgresql.TIMESTAMP(timezone=True), nullable=True), sa.Column("expires_at", postgresql.TIMESTAMP(timezone=True), nullable=True), sa.Column("log_prompts", sa.Boolean(), nullable=True), sa.Column( "metadata", postgresql.JSONB(), nullable=False, server_default=sa.text("'{}'::jsonb"), ), schema=SCHEMA, ) op.create_index( "idx_api_keys_prefix", "api_keys", ["prefix"], schema=SCHEMA, postgresql_where=sa.text("status = 'active'"), ) op.create_index("idx_api_keys_tenant", "api_keys", ["tenant_id"], schema=SCHEMA) # --- key_limits --- op.create_table( "key_limits", sa.Column( "key_id", postgresql.UUID(as_uuid=True), sa.ForeignKey(f"{SCHEMA}.api_keys.id", ondelete="CASCADE"), primary_key=True, ), sa.Column("rpm", sa.Integer(), nullable=True), sa.Column("tpm", sa.Integer(), nullable=True), sa.Column("concurrent", sa.Integer(), nullable=True), sa.Column("tokens_daily", sa.BigInteger(), nullable=True), sa.Column("tokens_monthly", sa.BigInteger(), nullable=True), sa.Column("tokens_total", sa.BigInteger(), nullable=True), sa.Column("allowed_models", postgresql.ARRAY(sa.Text()), nullable=True), sa.Column("allow_all_models", sa.Boolean(), nullable=True), schema=SCHEMA, ) # --- budget_usage --- op.create_table( "budget_usage", sa.Column( "key_id", postgresql.UUID(as_uuid=True), sa.ForeignKey(f"{SCHEMA}.api_keys.id", ondelete="CASCADE"), primary_key=True, nullable=False, ), sa.Column("period", _budget_period, primary_key=True, nullable=False), sa.Column( "period_start", postgresql.TIMESTAMP(timezone=True), primary_key=True, nullable=False, ), sa.Column("tokens_in", sa.BigInteger(), nullable=False, server_default=sa.text("0")), sa.Column("tokens_out", sa.BigInteger(), nullable=False, server_default=sa.text("0")), sa.Column("requests", sa.BigInteger(), nullable=False, server_default=sa.text("0")), schema=SCHEMA, ) op.create_index( "idx_budget_usage_period", "budget_usage", ["period", "period_start"], schema=SCHEMA, ) # --- audit_log --- op.create_table( "audit_log", sa.Column("id", sa.BigInteger(), sa.Identity(always=False), primary_key=True), sa.Column( "ts", postgresql.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("now()"), ), sa.Column("request_id", postgresql.UUID(as_uuid=True), nullable=False), sa.Column("tenant_id", postgresql.UUID(as_uuid=True), nullable=True), sa.Column("key_id", postgresql.UUID(as_uuid=True), nullable=True), sa.Column("key_prefix", sa.Text(), nullable=True), sa.Column("method", sa.Text(), nullable=False), sa.Column("path", sa.Text(), nullable=False), sa.Column("model", sa.Text(), nullable=True), sa.Column("tokens_in", sa.Integer(), nullable=True), sa.Column("tokens_out", sa.Integer(), nullable=True), sa.Column("latency_ms", sa.Integer(), nullable=True), sa.Column("status", sa.Integer(), nullable=False), sa.Column("client_ip", postgresql.INET(), nullable=True), sa.Column("user_agent", sa.Text(), nullable=True), sa.Column("error_code", sa.Text(), nullable=True), schema=SCHEMA, ) op.create_index("idx_audit_ts", "audit_log", ["ts"], schema=SCHEMA) op.create_index("idx_audit_tenant_ts", "audit_log", ["tenant_id", "ts"], schema=SCHEMA) op.create_index("idx_audit_key_ts", "audit_log", ["key_id", "ts"], schema=SCHEMA) # --- prompt_log --- op.create_table( "prompt_log", sa.Column("id", sa.BigInteger(), sa.Identity(always=False), primary_key=True), sa.Column( "audit_id", sa.BigInteger(), sa.ForeignKey(f"{SCHEMA}.audit_log.id", ondelete="CASCADE"), nullable=False, ), sa.Column( "ts", postgresql.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("now()"), ), sa.Column("key_id", postgresql.UUID(as_uuid=True), nullable=False), sa.Column("request_body", postgresql.JSONB(), nullable=False), sa.Column("response_text", sa.Text(), nullable=True), sa.Column("retention_until", postgresql.TIMESTAMP(timezone=True), nullable=False), schema=SCHEMA, ) op.create_index( "idx_prompt_log_retention", "prompt_log", ["retention_until"], schema=SCHEMA ) # --- revocations --- op.create_table( "revocations", sa.Column("id", sa.BigInteger(), sa.Identity(always=False), primary_key=True), sa.Column("key_id", postgresql.UUID(as_uuid=True), nullable=False), sa.Column( "ts", postgresql.TIMESTAMP(timezone=True), nullable=False, server_default=sa.text("now()"), ), sa.Column("reason", sa.Text(), nullable=True), sa.Column("processed_at", postgresql.TIMESTAMP(timezone=True), nullable=True), schema=SCHEMA, ) # --- NOTIFY trigger on revocation insert (SPEC §5) --- op.execute( """ CREATE OR REPLACE FUNCTION gateway.notify_key_revoked() RETURNS trigger AS $$ BEGIN PERFORM pg_notify('key_revoked', NEW.key_id::text); RETURN NEW; END; $$ LANGUAGE plpgsql; """ ) op.execute( """ CREATE TRIGGER trg_notify_key_revoked AFTER INSERT ON gateway.revocations FOR EACH ROW EXECUTE FUNCTION gateway.notify_key_revoked(); """ ) def downgrade() -> None: """Drop the entire ``gateway`` schema and its objects.""" op.execute("DROP TRIGGER IF EXISTS trg_notify_key_revoked ON gateway.revocations") op.execute("DROP FUNCTION IF EXISTS gateway.notify_key_revoked()") op.drop_index("idx_prompt_log_retention", table_name="prompt_log", schema=SCHEMA) op.drop_table("prompt_log", schema=SCHEMA) op.drop_index("idx_audit_key_ts", table_name="audit_log", schema=SCHEMA) op.drop_index("idx_audit_tenant_ts", table_name="audit_log", schema=SCHEMA) op.drop_index("idx_audit_ts", table_name="audit_log", schema=SCHEMA) op.drop_table("audit_log", schema=SCHEMA) op.drop_index("idx_budget_usage_period", table_name="budget_usage", schema=SCHEMA) op.drop_table("budget_usage", schema=SCHEMA) op.drop_table("key_limits", schema=SCHEMA) op.drop_index("idx_api_keys_tenant", table_name="api_keys", schema=SCHEMA) op.drop_index("idx_api_keys_prefix", table_name="api_keys", schema=SCHEMA) op.drop_table("api_keys", schema=SCHEMA) op.drop_table("tenant_limits", schema=SCHEMA) op.drop_table("tenants", schema=SCHEMA) op.execute("DROP TYPE IF EXISTS gateway.budget_period") op.execute("DROP TYPE IF EXISTS gateway.tenant_status") op.execute("DROP TYPE IF EXISTS gateway.key_status") op.execute(f'DROP SCHEMA IF EXISTS "{SCHEMA}"')