From 78116ec59dfc9a0d2d62115bf9f2335a2ac6e29a Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Mon, 18 May 2026 09:40:00 +0800 Subject: [PATCH] =?UTF-8?q?refactor(db):=20=E6=99=BA=E8=83=BD=E4=BD=93?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E4=BB=8E=E9=83=A8=E9=97=A8=E9=9A=94=E7=A6=BB?= =?UTF-8?q?=E6=94=B9=E4=B8=BA=E7=94=A8=E6=88=B7=E9=9A=94=E7=A6=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 修改 AgentConfig 模型 department_id -> uid - 添加数据库迁移脚本处理数据转换 - 更新 postgres manager 添加迁移逻辑 --- .../package/yuxi/storage/postgres/manager.py | 99 +++++++++++++ .../20260518_0001_agent_config_uid.py | 136 ++++++++++++++++++ .../yuxi/storage/postgres/models_business.py | 12 +- .../yuxi/storage/postgres/models_knowledge.py | 1 + 4 files changed, 242 insertions(+), 6 deletions(-) create mode 100644 backend/package/yuxi/storage/postgres/migrations/versions/20260518_0001_agent_config_uid.py diff --git a/backend/package/yuxi/storage/postgres/manager.py b/backend/package/yuxi/storage/postgres/manager.py index 7c708498..adabeab8 100644 --- a/backend/package/yuxi/storage/postgres/manager.py +++ b/backend/package/yuxi/storage/postgres/manager.py @@ -309,6 +309,105 @@ class PostgresManager(metaclass=SingletonMeta): "ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS is_pinned BOOLEAN NOT NULL DEFAULT FALSE", "ALTER TABLE IF EXISTS mcp_servers ADD COLUMN IF NOT EXISTS env JSONB", """ + ALTER TABLE IF EXISTS agent_configs + ADD COLUMN IF NOT EXISTS uid VARCHAR + """, + """ + UPDATE agent_configs ac + SET uid = u.uid + FROM users u + WHERE ac.uid IS NULL + AND ac.created_by ~ '^[0-9]+$' + AND u.id = ac.created_by::integer + """, + """ + UPDATE agent_configs ac + SET uid = u.uid + FROM users u + WHERE ac.uid IS NULL + AND ac.created_by = u.uid + """, + """ + DO $$ + BEGIN + IF EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'agent_configs' AND column_name = 'department_id' + ) THEN + EXECUTE $sql$ + UPDATE agent_configs ac + SET uid = ( + SELECT u.uid + FROM users u + WHERE u.department_id = ac.department_id + AND u.is_deleted = 0 + ORDER BY + CASE WHEN u.role = 'superadmin' THEN 0 WHEN u.role = 'admin' THEN 1 ELSE 2 END, + u.id ASC + LIMIT 1 + ) + WHERE ac.uid IS NULL + $sql$; + END IF; + END $$; + """, + "DELETE FROM agent_configs WHERE uid IS NULL", + "DROP INDEX IF EXISTS uq_agent_configs_department_agent_default", + "DROP INDEX IF EXISTS ix_agent_configs_department_id", + "ALTER TABLE IF EXISTS agent_configs DROP CONSTRAINT IF EXISTS uq_agent_configs_department_agent_name", + "ALTER TABLE IF EXISTS agent_configs DROP CONSTRAINT IF EXISTS agent_configs_department_id_fkey", + "ALTER TABLE IF EXISTS agent_configs DROP COLUMN IF EXISTS department_id", + "ALTER TABLE IF EXISTS agent_configs ALTER COLUMN uid SET NOT NULL", + """ + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (PARTITION BY uid, agent_id, name ORDER BY id) AS rn + FROM agent_configs + ) + UPDATE agent_configs ac + SET name = LEFT(ac.name, 90) || '-' || ac.id::text + FROM ranked + WHERE ac.id = ranked.id AND ranked.rn > 1 + """, + """ + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (PARTITION BY uid, agent_id ORDER BY id) AS rn + FROM agent_configs + WHERE is_default IS TRUE + ) + UPDATE agent_configs ac + SET is_default = FALSE + FROM ranked + WHERE ac.id = ranked.id AND ranked.rn > 1 + """, + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_constraint WHERE conname = 'agent_configs_uid_fkey' + ) THEN + ALTER TABLE agent_configs + ADD CONSTRAINT agent_configs_uid_fkey FOREIGN KEY (uid) REFERENCES users(uid); + END IF; + END $$; + """, + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_constraint WHERE conname = 'uq_agent_configs_uid_agent_name' + ) THEN + ALTER TABLE agent_configs + ADD CONSTRAINT uq_agent_configs_uid_agent_name UNIQUE (uid, agent_id, name); + END IF; + END $$; + """, + "CREATE INDEX IF NOT EXISTS ix_agent_configs_uid ON agent_configs(uid)", + """ + CREATE UNIQUE INDEX IF NOT EXISTS uq_agent_configs_uid_agent_default + ON agent_configs(uid, agent_id) + WHERE is_default IS TRUE + """, + """ CREATE TABLE IF NOT EXISTS model_providers ( id SERIAL PRIMARY KEY, provider_id VARCHAR(100) NOT NULL UNIQUE, diff --git a/backend/package/yuxi/storage/postgres/migrations/versions/20260518_0001_agent_config_uid.py b/backend/package/yuxi/storage/postgres/migrations/versions/20260518_0001_agent_config_uid.py new file mode 100644 index 00000000..cdaf894c --- /dev/null +++ b/backend/package/yuxi/storage/postgres/migrations/versions/20260518_0001_agent_config_uid.py @@ -0,0 +1,136 @@ +"""Move agent configs from department scope to user scope. + +Revision ID: 20260518_0001 +Revises: +Create Date: 2026-05-18 +""" + +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa + + +revision = "20260518_0001" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column("agent_configs", sa.Column("uid", sa.String(), nullable=True)) + + op.execute( + """ + UPDATE agent_configs ac + SET uid = u.uid + FROM users u + WHERE ac.uid IS NULL + AND ac.created_by ~ '^[0-9]+$' + AND u.id = ac.created_by::integer + """ + ) + op.execute( + """ + UPDATE agent_configs ac + SET uid = u.uid + FROM users u + WHERE ac.uid IS NULL + AND ac.created_by = u.uid + """ + ) + op.execute( + """ + UPDATE agent_configs ac + SET uid = ( + SELECT u.uid + FROM users u + WHERE u.department_id = ac.department_id + AND u.is_deleted = 0 + ORDER BY + CASE WHEN u.role = 'superadmin' THEN 0 WHEN u.role = 'admin' THEN 1 ELSE 2 END, + u.id ASC + LIMIT 1 + ) + WHERE ac.uid IS NULL + """ + ) + op.execute("DELETE FROM agent_configs WHERE uid IS NULL") + + op.execute("DROP INDEX IF EXISTS uq_agent_configs_department_agent_default") + op.execute("DROP INDEX IF EXISTS ix_agent_configs_department_id") + op.drop_constraint("uq_agent_configs_department_agent_name", "agent_configs", type_="unique") + op.drop_constraint("agent_configs_department_id_fkey", "agent_configs", type_="foreignkey") + op.drop_column("agent_configs", "department_id") + + op.alter_column("agent_configs", "uid", nullable=False) + op.execute( + """ + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (PARTITION BY uid, agent_id, name ORDER BY id) AS rn + FROM agent_configs + ) + UPDATE agent_configs ac + SET name = LEFT(ac.name, 90) || '-' || ac.id::text + FROM ranked + WHERE ac.id = ranked.id AND ranked.rn > 1 + """ + ) + op.execute( + """ + WITH ranked AS ( + SELECT id, ROW_NUMBER() OVER (PARTITION BY uid, agent_id ORDER BY id) AS rn + FROM agent_configs + WHERE is_default IS TRUE + ) + UPDATE agent_configs ac + SET is_default = FALSE + FROM ranked + WHERE ac.id = ranked.id AND ranked.rn > 1 + """ + ) + op.create_foreign_key("agent_configs_uid_fkey", "agent_configs", "users", ["uid"], ["uid"]) + op.create_unique_constraint("uq_agent_configs_uid_agent_name", "agent_configs", ["uid", "agent_id", "name"]) + op.create_index("ix_agent_configs_uid", "agent_configs", ["uid"]) + op.create_index( + "uq_agent_configs_uid_agent_default", + "agent_configs", + ["uid", "agent_id"], + unique=True, + postgresql_where=sa.text("is_default IS TRUE"), + ) + + +def downgrade() -> None: + op.add_column("agent_configs", sa.Column("department_id", sa.Integer(), nullable=True)) + op.execute( + """ + UPDATE agent_configs ac + SET department_id = u.department_id + FROM users u + WHERE ac.uid = u.uid + """ + ) + op.execute("DELETE FROM agent_configs WHERE department_id IS NULL") + + op.execute("DROP INDEX IF EXISTS uq_agent_configs_uid_agent_default") + op.execute("DROP INDEX IF EXISTS ix_agent_configs_uid") + op.drop_constraint("uq_agent_configs_uid_agent_name", "agent_configs", type_="unique") + op.drop_constraint("agent_configs_uid_fkey", "agent_configs", type_="foreignkey") + op.drop_column("agent_configs", "uid") + + op.alter_column("agent_configs", "department_id", nullable=False) + op.create_foreign_key("agent_configs_department_id_fkey", "agent_configs", "departments", ["department_id"], ["id"]) + op.create_unique_constraint( + "uq_agent_configs_department_agent_name", + "agent_configs", + ["department_id", "agent_id", "name"], + ) + op.create_index("ix_agent_configs_department_id", "agent_configs", ["department_id"]) + op.create_index( + "uq_agent_configs_department_agent_default", + "agent_configs", + ["department_id", "agent_id"], + unique=True, + postgresql_where=sa.text("is_default IS TRUE"), + ) diff --git a/backend/package/yuxi/storage/postgres/models_business.py b/backend/package/yuxi/storage/postgres/models_business.py index 1ac64cff..96ba528d 100644 --- a/backend/package/yuxi/storage/postgres/models_business.py +++ b/backend/package/yuxi/storage/postgres/models_business.py @@ -131,12 +131,12 @@ class User(Base): class AgentConfig(Base): - """智能体配置（按部门共享，多份可切换）""" + """智能体配置（按用户隔离，多份可切换）""" __tablename__ = "agent_configs" id = Column(Integer, primary_key=True, autoincrement=True) - department_id = Column(Integer, ForeignKey("departments.id"), nullable=False, index=True) + uid = Column(String, ForeignKey("users.uid"), nullable=False, index=True) agent_id = Column(String(64), nullable=False, index=True) name = Column(String(100), nullable=False) @@ -155,10 +155,10 @@ class AgentConfig(Base): updated_at = Column(DateTime, default=utc_now_naive, onupdate=utc_now_naive) __table_args__ = ( - UniqueConstraint("department_id", "agent_id", "name", name="uq_agent_configs_department_agent_name"), + UniqueConstraint("uid", "agent_id", "name", name="uq_agent_configs_uid_agent_name"), Index( - "uq_agent_configs_department_agent_default", - "department_id", + "uq_agent_configs_uid_agent_default", + "uid", "agent_id", unique=True, postgresql_where=is_default.is_(True), @@ -168,7 +168,7 @@ class AgentConfig(Base): def to_dict(self) -> dict[str, Any]: return { "id": self.id, - "department_id": self.department_id, + "uid": self.uid, "agent_id": self.agent_id, "name": self.name, "description": self.description, diff --git a/backend/package/yuxi/storage/postgres/models_knowledge.py b/backend/package/yuxi/storage/postgres/models_knowledge.py index dcd2462b..29a86824 100644 --- a/backend/package/yuxi/storage/postgres/models_knowledge.py +++ b/backend/package/yuxi/storage/postgres/models_knowledge.py @@ -39,6 +39,7 @@ class KnowledgeBase(Base): share_config = Column(JSON_VALUE) mindmap = Column(JSON_VALUE) sample_questions = Column(JSON_VALUE) + created_by = Column(String(64)) created_at = Column(DateTime(timezone=True), default=utc_now_naive) updated_at = Column(DateTime(timezone=True), default=utc_now_naive, onupdate=utc_now_naive)