feat(migration): add user_agent_catalog table and update trigger
This commit is contained in:
@@ -0,0 +1,173 @@
|
|||||||
|
"""add_user_agent_catalog
|
||||||
|
|
||||||
|
Revision ID: 50ae013ce530
|
||||||
|
Revises: 202602270006
|
||||||
|
Create Date: 2026-03-02 15:34:25.995336
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from sqlalchemy.dialects import postgresql
|
||||||
|
|
||||||
|
|
||||||
|
revision: str = "50ae013ce530"
|
||||||
|
down_revision: Union[str, Sequence[str], None] = "202602270006"
|
||||||
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.create_table(
|
||||||
|
"user_agent_catalog",
|
||||||
|
sa.Column("agent_type", sa.String(20), nullable=False),
|
||||||
|
sa.Column("llm_id", sa.UUID(), nullable=False),
|
||||||
|
sa.Column("status", sa.String(20), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
"config",
|
||||||
|
postgresql.JSONB(astext_type=sa.Text()),
|
||||||
|
server_default="{}",
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column(
|
||||||
|
"created_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
server_default=sa.text("now()"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column(
|
||||||
|
"updated_at",
|
||||||
|
sa.DateTime(timezone=True),
|
||||||
|
server_default=sa.text("now()"),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.PrimaryKeyConstraint("agent_type"),
|
||||||
|
sa.ForeignKeyConstraint(
|
||||||
|
["llm_id"],
|
||||||
|
["llms.id"],
|
||||||
|
name="fk_user_agent_catalog_llm_id",
|
||||||
|
ondelete="RESTRICT",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
op.execute(
|
||||||
|
"ALTER TABLE user_agent_catalog "
|
||||||
|
"ADD CONSTRAINT chk_user_agent_catalog_status "
|
||||||
|
"CHECK (status IN ('active', 'paused', 'migrating'))"
|
||||||
|
)
|
||||||
|
|
||||||
|
_enable_rls("user_agent_catalog")
|
||||||
|
|
||||||
|
op.execute("""
|
||||||
|
CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
|
||||||
|
RETURNS trigger
|
||||||
|
LANGUAGE plpgsql
|
||||||
|
SECURITY DEFINER
|
||||||
|
SET search_path = public
|
||||||
|
AS $$
|
||||||
|
BEGIN
|
||||||
|
INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
|
||||||
|
VALUES (
|
||||||
|
NEW.id,
|
||||||
|
COALESCE(
|
||||||
|
NEW.raw_user_meta_data ->> 'username',
|
||||||
|
split_part(NEW.email, '@', 1),
|
||||||
|
'user_' || substring(NEW.id::text, 1, 8)
|
||||||
|
),
|
||||||
|
NULL,
|
||||||
|
NULL,
|
||||||
|
'{}'::jsonb,
|
||||||
|
now(),
|
||||||
|
now()
|
||||||
|
)
|
||||||
|
ON CONFLICT (id) DO NOTHING;
|
||||||
|
|
||||||
|
INSERT INTO public.user_agents (id, user_id, llm_id, agent_type, config, status, created_by, updated_by)
|
||||||
|
SELECT
|
||||||
|
gen_random_uuid(),
|
||||||
|
NEW.id,
|
||||||
|
uac.llm_id,
|
||||||
|
uac.agent_type,
|
||||||
|
uac.config,
|
||||||
|
uac.status,
|
||||||
|
NEW.id,
|
||||||
|
NEW.id
|
||||||
|
FROM public.user_agent_catalog uac;
|
||||||
|
|
||||||
|
RETURN NEW;
|
||||||
|
END;
|
||||||
|
$$;
|
||||||
|
""")
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.execute("""
|
||||||
|
CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
|
||||||
|
RETURNS trigger
|
||||||
|
LANGUAGE plpgsql
|
||||||
|
SECURITY DEFINER
|
||||||
|
SET search_path = public
|
||||||
|
AS $$
|
||||||
|
BEGIN
|
||||||
|
INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
|
||||||
|
VALUES (
|
||||||
|
NEW.id,
|
||||||
|
COALESCE(
|
||||||
|
NEW.raw_user_meta_data ->> 'username',
|
||||||
|
split_part(NEW.email, '@', 1),
|
||||||
|
'user_' || substring(NEW.id::text, 1, 8)
|
||||||
|
),
|
||||||
|
NULL,
|
||||||
|
NULL,
|
||||||
|
'{}'::jsonb,
|
||||||
|
now(),
|
||||||
|
now()
|
||||||
|
)
|
||||||
|
ON CONFLICT (id) DO NOTHING;
|
||||||
|
|
||||||
|
RETURN NEW;
|
||||||
|
END;
|
||||||
|
$$;
|
||||||
|
""")
|
||||||
|
|
||||||
|
_drop_rls("user_agent_catalog")
|
||||||
|
op.drop_constraint(
|
||||||
|
"chk_user_agent_catalog_status", "user_agent_catalog", type_="check"
|
||||||
|
)
|
||||||
|
op.drop_constraint(
|
||||||
|
"fk_user_agent_catalog_llm_id", "user_agent_catalog", type_="foreignkey"
|
||||||
|
)
|
||||||
|
op.drop_table("user_agent_catalog")
|
||||||
|
|
||||||
|
|
||||||
|
def _enable_rls(table_name: str) -> None:
|
||||||
|
for role in ["anon", "authenticated"]:
|
||||||
|
for action in ["select", "insert", "update", "delete"]:
|
||||||
|
op.execute(
|
||||||
|
f"DROP POLICY IF EXISTS {role}_{action}_{table_name} ON {table_name}"
|
||||||
|
)
|
||||||
|
op.execute(f"ALTER TABLE {table_name} ENABLE ROW LEVEL SECURITY")
|
||||||
|
for role in ["anon", "authenticated"]:
|
||||||
|
op.execute(
|
||||||
|
f"CREATE POLICY {role}_select_{table_name} ON {table_name} FOR SELECT TO {role} USING (false)"
|
||||||
|
)
|
||||||
|
op.execute(
|
||||||
|
f"CREATE POLICY {role}_insert_{table_name} ON {table_name} FOR INSERT TO {role} WITH CHECK (false)"
|
||||||
|
)
|
||||||
|
op.execute(
|
||||||
|
f"CREATE POLICY {role}_update_{table_name} ON {table_name} FOR UPDATE TO {role} USING (false) WITH CHECK (false)"
|
||||||
|
)
|
||||||
|
op.execute(
|
||||||
|
f"CREATE POLICY {role}_delete_{table_name} ON {table_name} FOR DELETE TO {role} USING (false)"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _drop_rls(table_name: str) -> None:
|
||||||
|
for role in ["anon", "authenticated"]:
|
||||||
|
op.execute(f"DROP POLICY IF EXISTS {role}_delete_{table_name} ON {table_name}")
|
||||||
|
op.execute(f"DROP POLICY IF EXISTS {role}_update_{table_name} ON {table_name}")
|
||||||
|
op.execute(f"DROP POLICY IF EXISTS {role}_insert_{table_name} ON {table_name}")
|
||||||
|
op.execute(f"DROP POLICY IF EXISTS {role}_select_{table_name} ON {table_name}")
|
||||||
|
op.execute(f"ALTER TABLE {table_name} DISABLE ROW LEVEL SECURITY")
|
||||||
Reference in New Issue
Block a user