From 27b09ce9c0944f9938e754a626e713d8eb605da9 Mon Sep 17 00:00:00 2001 From: qzl Date: Mon, 2 Mar 2026 15:36:25 +0800 Subject: [PATCH] feat(migration): add user_agent_catalog table and update trigger --- .../50ae013ce530_add_user_agent_catalog.py | 173 ++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 backend/alembic/versions/50ae013ce530_add_user_agent_catalog.py diff --git a/backend/alembic/versions/50ae013ce530_add_user_agent_catalog.py b/backend/alembic/versions/50ae013ce530_add_user_agent_catalog.py new file mode 100644 index 0000000..b5b106d --- /dev/null +++ b/backend/alembic/versions/50ae013ce530_add_user_agent_catalog.py @@ -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")