Skip to content

Commit 85aed8e

Browse files
author
Francisco
committed
fix: enhance migrations with deferred FK creation, safe column alterations, and improved table guards
1 parent 225f155 commit 85aed8e

5 files changed

Lines changed: 192 additions & 154 deletions

‎migrations/versions/26a927cbe516_remove_thread_vector_store_relationship.py‎

Lines changed: 65 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -12,33 +12,30 @@
1212
from alembic import op
1313
from sqlalchemy.dialects import mysql
1414

15-
# Import SafeDDL helpers
16-
from migrations.utils.safe_ddl import has_table, safe_alter_column
15+
from migrations.utils.safe_ddl import (create_fk_if_not_exists, has_table,
16+
safe_alter_column)
1717

1818
# revision identifiers, used by Alembic.
19-
revision: str = '26a927cbe516'
20-
down_revision: Union[str, None] = '3a42e4f129e4'
19+
revision: str = "26a927cbe516"
20+
down_revision: Union[str, None] = "3a42e4f129e4"
2121
branch_labels: Union[str, Sequence[str], None] = None
2222
depends_on: Union[str, Sequence[str], None] = None
2323

2424

2525
def upgrade() -> None:
2626
"""Upgrade schema."""
27-
# Check if tables exist before attempting to drop them
28-
if has_table('thread_vector_stores'):
29-
op.drop_table('thread_vector_stores')
27+
if has_table("thread_vector_stores"):
28+
op.drop_table("thread_vector_stores")
3029

31-
if has_table('vector_store_assistants'):
32-
op.drop_table('vector_store_assistants')
30+
if has_table("vector_store_assistants"):
31+
op.drop_table("vector_store_assistants")
3332

34-
# Safe column alterations
3533
safe_alter_column(
36-
table_name='messages', column_name='content', existing_type=mysql.TEXT(), nullable=False
34+
table_name="messages", column_name="content", existing_type=mysql.TEXT(), nullable=False
3735
)
38-
3936
safe_alter_column(
40-
table_name='messages',
41-
column_name='reasoning',
37+
table_name="messages",
38+
column_name="reasoning",
4239
existing_type=mysql.LONGTEXT(),
4340
type_=sa.Text(length=4294967295),
4441
existing_comment="Stores the internal 'thinking' or reasoning tokens from the model.",
@@ -48,53 +45,71 @@ def upgrade() -> None:
4845

4946
def downgrade() -> None:
5047
"""Downgrade schema."""
51-
# Safe column alterations for reversion
5248
safe_alter_column(
53-
table_name='messages',
54-
column_name='reasoning',
49+
table_name="messages",
50+
column_name="reasoning",
5551
existing_type=sa.Text(length=4294967295),
5652
type_=mysql.LONGTEXT(),
5753
existing_comment="Stores the internal 'thinking' or reasoning tokens from the model.",
5854
existing_nullable=True,
5955
)
60-
6156
safe_alter_column(
62-
table_name='messages', column_name='content', existing_type=mysql.TEXT(), nullable=True
57+
table_name="messages", column_name="content", existing_type=mysql.TEXT(), nullable=True
6358
)
6459

65-
# Check if tables are missing before attempting to create them
66-
if not has_table('vector_store_assistants'):
60+
# Recreate join tables without inline FKs.
61+
# 'assistants', 'vector_stores', and 'threads' are base tables that may not
62+
# exist yet on a fresh container at this point in the downgrade chain.
63+
# FKs are added as separate deferred steps via create_fk_if_not_exists below,
64+
# which checks has_table on both source and referent before acting.
65+
if not has_table("vector_store_assistants"):
6766
op.create_table(
68-
'vector_store_assistants',
69-
sa.Column('vector_store_id', mysql.VARCHAR(length=64), nullable=False),
70-
sa.Column('assistant_id', mysql.VARCHAR(length=64), nullable=False),
71-
sa.ForeignKeyConstraint(
72-
['assistant_id'], ['assistants.id'], name=op.f('vector_store_assistants_ibfk_2')
73-
),
74-
sa.ForeignKeyConstraint(
75-
['vector_store_id'],
76-
['vector_stores.id'],
77-
name=op.f('vector_store_assistants_ibfk_1'),
78-
),
79-
sa.PrimaryKeyConstraint('vector_store_id', 'assistant_id'),
80-
mysql_collate='utf8mb4_0900_ai_ci',
81-
mysql_default_charset='utf8mb4',
82-
mysql_engine='InnoDB',
67+
"vector_store_assistants",
68+
sa.Column("vector_store_id", mysql.VARCHAR(length=64), nullable=False),
69+
sa.Column("assistant_id", mysql.VARCHAR(length=64), nullable=False),
70+
sa.PrimaryKeyConstraint("vector_store_id", "assistant_id"),
71+
mysql_collate="utf8mb4_0900_ai_ci",
72+
mysql_default_charset="utf8mb4",
73+
mysql_engine="InnoDB",
8374
)
8475

85-
if not has_table('thread_vector_stores'):
76+
create_fk_if_not_exists(
77+
"vector_store_assistants_ibfk_1",
78+
"vector_store_assistants",
79+
"vector_stores",
80+
["vector_store_id"],
81+
["id"],
82+
)
83+
create_fk_if_not_exists(
84+
"vector_store_assistants_ibfk_2",
85+
"vector_store_assistants",
86+
"assistants",
87+
["assistant_id"],
88+
["id"],
89+
)
90+
91+
if not has_table("thread_vector_stores"):
8692
op.create_table(
87-
'thread_vector_stores',
88-
sa.Column('thread_id', mysql.VARCHAR(length=64), nullable=False),
89-
sa.Column('vector_store_id', mysql.VARCHAR(length=64), nullable=False),
90-
sa.ForeignKeyConstraint(
91-
['thread_id'], ['threads.id'], name=op.f('thread_vector_stores_ibfk_1')
92-
),
93-
sa.ForeignKeyConstraint(
94-
['vector_store_id'], ['vector_stores.id'], name=op.f('thread_vector_stores_ibfk_2')
95-
),
96-
sa.PrimaryKeyConstraint('thread_id', 'vector_store_id'),
97-
mysql_collate='utf8mb4_0900_ai_ci',
98-
mysql_default_charset='utf8mb4',
99-
mysql_engine='InnoDB',
93+
"thread_vector_stores",
94+
sa.Column("thread_id", mysql.VARCHAR(length=64), nullable=False),
95+
sa.Column("vector_store_id", mysql.VARCHAR(length=64), nullable=False),
96+
sa.PrimaryKeyConstraint("thread_id", "vector_store_id"),
97+
mysql_collate="utf8mb4_0900_ai_ci",
98+
mysql_default_charset="utf8mb4",
99+
mysql_engine="InnoDB",
100100
)
101+
102+
create_fk_if_not_exists(
103+
"thread_vector_stores_ibfk_1",
104+
"thread_vector_stores",
105+
"threads",
106+
["thread_id"],
107+
["id"],
108+
)
109+
create_fk_if_not_exists(
110+
"thread_vector_stores_ibfk_2",
111+
"thread_vector_stores",
112+
"vector_stores",
113+
["vector_store_id"],
114+
["id"],
115+
)

‎migrations/versions/3a42e4f129e4_remove_assistant_vector_store_.py‎

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -12,31 +12,32 @@
1212
from alembic import op
1313
from sqlalchemy.dialects import mysql
1414

15-
from migrations.utils.safe_ddl import has_table, safe_alter_column
15+
from migrations.utils.safe_ddl import (create_fk_if_not_exists, has_table,
16+
safe_alter_column)
1617

1718
# revision identifiers, used by Alembic.
18-
revision: str = '3a42e4f129e4'
19-
down_revision: Union[str, None] = '1c9784351972'
19+
revision: str = "3a42e4f129e4"
20+
down_revision: Union[str, None] = "1c9784351972"
2021
branch_labels: Union[str, Sequence[str], None] = None
2122
depends_on: Union[str, Sequence[str], None] = None
2223

2324
_JOIN_TABLE = "vector_store_assistants"
2425

2526

2627
def upgrade() -> None:
27-
# ── 1. Drop the M2M join table if it still exists ─────────────────────
28+
# ── 1. Drop the M2M join table if it still exists ──────────────────────
2829
if has_table(_JOIN_TABLE):
2930
op.drop_table(_JOIN_TABLE)
3031

31-
# ── 2. messages.content → NOT NULL ────────────────────────────────────
32+
# ── 2. messages.content → NOT NULL ─────────────────────────────────────
3233
safe_alter_column(
3334
"messages",
3435
"content",
3536
existing_type=mysql.TEXT(),
3637
nullable=False,
3738
)
3839

39-
# ── 3. messages.reasoning → LONGTEXT (explicit length form) ───────────
40+
# ── 3. messages.reasoning → LONGTEXT (explicit length form) ────────────
4041
safe_alter_column(
4142
"messages",
4243
"reasoning",
@@ -48,7 +49,7 @@ def upgrade() -> None:
4849

4950

5051
def downgrade() -> None:
51-
# ── 3. Restore messages.reasoning type ────────────────────────────────
52+
# ── 3. Restore messages.reasoning type ─────────────────────────────────
5253
safe_alter_column(
5354
"messages",
5455
"reasoning",
@@ -58,30 +59,39 @@ def downgrade() -> None:
5859
existing_nullable=True,
5960
)
6061

61-
# ── 2. messages.content → nullable again ──────────────────────────────
62+
# ── 2. messages.content → nullable again ───────────────────────────────
6263
safe_alter_column(
6364
"messages",
6465
"content",
6566
existing_type=mysql.TEXT(),
6667
nullable=True,
6768
)
6869

69-
# ── 1. Recreate the join table (schema only, data is gone) ────────────
70+
# ── 1. Recreate the join table without inline FKs ──────────────────────
71+
# 'vector_stores' and 'assistants' are base tables that may not exist yet
72+
# on a fresh container at this point in the downgrade chain. Inline FK
73+
# declarations would cause MySQL error 1824. FKs are added separately via
74+
# create_fk_if_not_exists, which checks has_table on both tables first.
7075
if not has_table(_JOIN_TABLE):
7176
op.create_table(
7277
_JOIN_TABLE,
73-
sa.Column(
74-
"vector_store_id",
75-
sa.String(64),
76-
sa.ForeignKey("vector_stores.id", ondelete="CASCADE"),
77-
primary_key=True,
78-
nullable=False,
79-
),
80-
sa.Column(
81-
"assistant_id",
82-
sa.String(64),
83-
sa.ForeignKey("assistants.id", ondelete="CASCADE"),
84-
primary_key=True,
85-
nullable=False,
86-
),
78+
sa.Column("vector_store_id", sa.String(64), primary_key=True, nullable=False),
79+
sa.Column("assistant_id", sa.String(64), primary_key=True, nullable=False),
8780
)
81+
82+
create_fk_if_not_exists(
83+
"vector_store_assistants_ibfk_1",
84+
_JOIN_TABLE,
85+
"vector_stores",
86+
["vector_store_id"],
87+
["id"],
88+
ondelete="CASCADE",
89+
)
90+
create_fk_if_not_exists(
91+
"vector_store_assistants_ibfk_2",
92+
_JOIN_TABLE,
93+
"assistants",
94+
["assistant_id"],
95+
["id"],
96+
ondelete="CASCADE",
97+
)

0 commit comments

Comments
 (0)