feat(migration): add user_agent_catalog table and update trigger

2026-03-02 15:36:25 +08:00
parent 97626ec2c4
commit 27b09ce9c0
1 changed files with 173 additions and 0 deletions
@@ -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")