fix: make migration database-agnostic for SQLite and PostgreSQL

This commit is contained in:
2026-05-22 20:48:21 +00:00
parent 7cc720786e
commit 0901b1e832
@@ -10,6 +10,7 @@ from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
from sqlalchemy import inspect
# revision identifiers, used by Alembic.
revision: str = "0015_interface_type_requires_port"
@@ -18,26 +19,57 @@ branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def _get_dialect() -> str:
"""Get the current database dialect name."""
conn = op.get_bind()
return conn.dialect.name
def upgrade() -> None:
dialect = _get_dialect()
# Add new columns
op.add_column('tool_types', sa.Column('interface_type', sa.String(20), nullable=True))
op.add_column('tool_types', sa.Column('requires_port', sa.Boolean(), nullable=False, server_default='true'))
# Migrate data: take first element from interfaces JSON array
op.execute("""
UPDATE tool_types
SET interface_type = COALESCE(
(SELECT value->>0 FROM jsonb_array_elements_text(interfaces) AS value LIMIT 1),
'web'
),
requires_port = CASE
WHEN COALESCE(
(SELECT value->>0 FROM jsonb_array_elements_text(interfaces) AS value LIMIT 1),
if dialect == 'postgresql':
op.execute("""
UPDATE tool_types
SET interface_type = COALESCE(
(SELECT elem FROM jsonb_array_elements_text(interfaces::jsonb) AS elem LIMIT 1),
'web'
) = 'web' THEN true
ELSE false
END
""")
),
requires_port = CASE
WHEN COALESCE(
(SELECT elem FROM jsonb_array_elements_text(interfaces::jsonb) AS elem LIMIT 1),
'web'
) = 'web' THEN true
ELSE false
END
""")
else:
# SQLite: interfaces is stored as JSON text, extract first array element
op.execute("""
UPDATE tool_types
SET interface_type = COALESCE(
(SELECT json_extract(value, '$[0]')
FROM json_each(interfaces) AS value
WHERE json_valid(interfaces)
LIMIT 1),
'web'
),
requires_port = CASE
WHEN COALESCE(
(SELECT json_extract(value, '$[0]')
FROM json_each(interfaces) AS value
WHERE json_valid(interfaces)
LIMIT 1),
'web'
) = 'web' THEN true
ELSE false
END
""")
# Make interface_type non-nullable after data migration
op.alter_column('tool_types', 'interface_type', nullable=False)
@@ -45,22 +77,33 @@ def upgrade() -> None:
# Drop old interfaces column
op.drop_column('tool_types', 'interfaces')
# Add CHECK constraint for interface_type
# Add CHECK constraint for interface_type (only on PostgreSQL; SQLite supports it too)
op.create_check_constraint('chk_interface_type', 'tool_types', sa.text("interface_type IN ('web', 'terminal')"))
def downgrade() -> None:
dialect = _get_dialect()
# Drop CHECK constraint
op.drop_constraint('chk_interface_type', 'tool_types', type_='check')
# Add back interfaces column
op.add_column('tool_types', sa.Column('interfaces', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='["web"]'))
# Migrate data back: wrap interface_type in array
op.execute("""
UPDATE tool_types
SET interfaces = jsonb_build_array(interface_type)
""")
if dialect == 'postgresql':
op.add_column('tool_types', sa.Column('interfaces', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='["web"]'))
# Migrate data back: wrap interface_type in array
op.execute("""
UPDATE tool_types
SET interfaces = jsonb_build_array(interface_type)
""")
else:
op.add_column('tool_types', sa.Column('interfaces', sa.JSON(), nullable=False, server_default='["web"]'))
# Migrate data back: wrap interface_type in array for SQLite
op.execute("""
UPDATE tool_types
SET interfaces = json_array(interface_type)
""")
# Drop new columns
op.drop_column('tool_types', 'requires_port')