diff --git a/backend/open_webui/migrations/versions/461111b60977_add_missing_primary_keys_to_legacy_.py b/backend/open_webui/migrations/versions/461111b60977_add_missing_primary_keys_to_legacy_.py index c061e1febc..67b737c530 100644 --- a/backend/open_webui/migrations/versions/461111b60977_add_missing_primary_keys_to_legacy_.py +++ b/backend/open_webui/migrations/versions/461111b60977_add_missing_primary_keys_to_legacy_.py @@ -19,10 +19,12 @@ depends_on: Union[str, Sequence[str], None] = None # Tables bootstrapped by the old Peewee migration layer that may have # UNIQUE(id) but no PRIMARY KEY constraint. Fresh Alembic installs # already have correct PKs from 7e5b5dc7342b_init.py. -LEGACY_TABLES = [ - 'auth', 'chat', 'chatidtag', 'document', 'file', - 'function', 'memory', 'model', 'prompt', 'tag', 'tool', 'user', -] +# 'tag' uses a composite PK since the same tag name can exist for multiple users. +LEGACY_TABLES = { + 'auth': ['id'], 'chat': ['id'], 'chatidtag': ['id'], 'document': ['id'], + 'file': ['id'], 'function': ['id'], 'memory': ['id'], 'model': ['id'], + 'prompt': ['id'], 'tag': ['id', 'user_id'], 'tool': ['id'], 'user': ['id'], +} def upgrade() -> None: @@ -30,30 +32,31 @@ def upgrade() -> None: inspector = sa.inspect(conn) existing_tables = set(inspector.get_table_names()) - for table_name in LEGACY_TABLES: + for table_name, pk_columns in LEGACY_TABLES.items(): if table_name not in existing_tables: continue pk = inspector.get_pk_constraint(table_name) pk_cols = pk.get('constrained_columns', []) - # Already has a proper PK on 'id' — nothing to do - if pk_cols == ['id']: + # Already has the correct PK — nothing to do + if sorted(pk_cols) == sorted(pk_columns): continue - # Check that an 'id' column actually exists + # Check that all PK columns exist columns = {c['name'] for c in inspector.get_columns(table_name)} - if 'id' not in columns: + if not all(c in columns for c in pk_columns): continue - print(f"Promoting UNIQUE(id) → PRIMARY KEY for '{table_name}'") + print(f"Promoting UNIQUE(id) -> PRIMARY KEY({', '.join(pk_columns)}) for '{table_name}'") + conn.execute(sa.text(f'DROP TABLE IF EXISTS _alembic_tmp_{table_name}')) with op.batch_alter_table(table_name) as batch_op: # Drop existing PK if any (e.g. on wrong column) if pk_cols and pk.get('name'): batch_op.drop_constraint(pk['name'], type_='primary') - batch_op.create_primary_key(f'pk_{table_name}_id', ['id']) + batch_op.create_primary_key(f'pk_{table_name}', pk_columns) def downgrade() -> None: