fix: make migration database-agnostic for SQLite and PostgreSQL
This commit is contained in:
@@ -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')
|
||||
|
||||
Reference in New Issue
Block a user