|
19 | 19 | # Tables bootstrapped by the old Peewee migration layer that may have |
20 | 20 | # UNIQUE(id) but no PRIMARY KEY constraint. Fresh Alembic installs |
21 | 21 | # already have correct PKs from 7e5b5dc7342b_init.py. |
22 | | -LEGACY_TABLES = [ |
23 | | - 'auth', 'chat', 'chatidtag', 'document', 'file', |
24 | | - 'function', 'memory', 'model', 'prompt', 'tag', 'tool', 'user', |
25 | | -] |
| 22 | +# 'tag' uses a composite PK since the same tag name can exist for multiple users. |
| 23 | +LEGACY_TABLES = { |
| 24 | + 'auth': ['id'], 'chat': ['id'], 'chatidtag': ['id'], 'document': ['id'], |
| 25 | + 'file': ['id'], 'function': ['id'], 'memory': ['id'], 'model': ['id'], |
| 26 | + 'prompt': ['id'], 'tag': ['id', 'user_id'], 'tool': ['id'], 'user': ['id'], |
| 27 | +} |
26 | 28 |
|
27 | 29 |
|
28 | 30 | def upgrade() -> None: |
29 | 31 | conn = op.get_bind() |
30 | 32 | inspector = sa.inspect(conn) |
31 | 33 | existing_tables = set(inspector.get_table_names()) |
32 | 34 |
|
33 | | - for table_name in LEGACY_TABLES: |
| 35 | + for table_name, pk_columns in LEGACY_TABLES.items(): |
34 | 36 | if table_name not in existing_tables: |
35 | 37 | continue |
36 | 38 |
|
37 | 39 | pk = inspector.get_pk_constraint(table_name) |
38 | 40 | pk_cols = pk.get('constrained_columns', []) |
39 | 41 |
|
40 | | - # Already has a proper PK on 'id' — nothing to do |
41 | | - if pk_cols == ['id']: |
| 42 | + # Already has the correct PK — nothing to do |
| 43 | + if sorted(pk_cols) == sorted(pk_columns): |
42 | 44 | continue |
43 | 45 |
|
44 | | - # Check that an 'id' column actually exists |
| 46 | + # Check that all PK columns exist |
45 | 47 | columns = {c['name'] for c in inspector.get_columns(table_name)} |
46 | | - if 'id' not in columns: |
| 48 | + if not all(c in columns for c in pk_columns): |
47 | 49 | continue |
48 | 50 |
|
49 | | - print(f"Promoting UNIQUE(id) → PRIMARY KEY for '{table_name}'") |
| 51 | + print(f"Promoting UNIQUE(id) -> PRIMARY KEY({', '.join(pk_columns)}) for '{table_name}'") |
50 | 52 |
|
| 53 | + conn.execute(sa.text(f'DROP TABLE IF EXISTS _alembic_tmp_{table_name}')) |
51 | 54 | with op.batch_alter_table(table_name) as batch_op: |
52 | 55 | # Drop existing PK if any (e.g. on wrong column) |
53 | 56 | if pk_cols and pk.get('name'): |
54 | 57 | batch_op.drop_constraint(pk['name'], type_='primary') |
55 | 58 |
|
56 | | - batch_op.create_primary_key(f'pk_{table_name}_id', ['id']) |
| 59 | + batch_op.create_primary_key(f'pk_{table_name}', pk_columns) |
57 | 60 |
|
58 | 61 |
|
59 | 62 | def downgrade() -> None: |
|
0 commit comments