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
|
from alembic import op
|
||||||
import sqlalchemy as sa
|
import sqlalchemy as sa
|
||||||
from sqlalchemy.dialects import postgresql
|
from sqlalchemy.dialects import postgresql
|
||||||
|
from sqlalchemy import inspect
|
||||||
|
|
||||||
# revision identifiers, used by Alembic.
|
# revision identifiers, used by Alembic.
|
||||||
revision: str = "0015_interface_type_requires_port"
|
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
|
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:
|
def upgrade() -> None:
|
||||||
|
dialect = _get_dialect()
|
||||||
|
|
||||||
# Add new columns
|
# Add new columns
|
||||||
op.add_column('tool_types', sa.Column('interface_type', sa.String(20), nullable=True))
|
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'))
|
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
|
# Migrate data: take first element from interfaces JSON array
|
||||||
op.execute("""
|
if dialect == 'postgresql':
|
||||||
UPDATE tool_types
|
op.execute("""
|
||||||
SET interface_type = COALESCE(
|
UPDATE tool_types
|
||||||
(SELECT value->>0 FROM jsonb_array_elements_text(interfaces) AS value LIMIT 1),
|
SET interface_type = COALESCE(
|
||||||
'web'
|
(SELECT elem FROM jsonb_array_elements_text(interfaces::jsonb) AS elem LIMIT 1),
|
||||||
),
|
|
||||||
requires_port = CASE
|
|
||||||
WHEN COALESCE(
|
|
||||||
(SELECT value->>0 FROM jsonb_array_elements_text(interfaces) AS value LIMIT 1),
|
|
||||||
'web'
|
'web'
|
||||||
) = 'web' THEN true
|
),
|
||||||
ELSE false
|
requires_port = CASE
|
||||||
END
|
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
|
# Make interface_type non-nullable after data migration
|
||||||
op.alter_column('tool_types', 'interface_type', nullable=False)
|
op.alter_column('tool_types', 'interface_type', nullable=False)
|
||||||
@@ -45,22 +77,33 @@ def upgrade() -> None:
|
|||||||
# Drop old interfaces column
|
# Drop old interfaces column
|
||||||
op.drop_column('tool_types', 'interfaces')
|
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')"))
|
op.create_check_constraint('chk_interface_type', 'tool_types', sa.text("interface_type IN ('web', 'terminal')"))
|
||||||
|
|
||||||
|
|
||||||
def downgrade() -> None:
|
def downgrade() -> None:
|
||||||
|
dialect = _get_dialect()
|
||||||
|
|
||||||
# Drop CHECK constraint
|
# Drop CHECK constraint
|
||||||
op.drop_constraint('chk_interface_type', 'tool_types', type_='check')
|
op.drop_constraint('chk_interface_type', 'tool_types', type_='check')
|
||||||
|
|
||||||
# Add back interfaces column
|
# Add back interfaces column
|
||||||
op.add_column('tool_types', sa.Column('interfaces', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='["web"]'))
|
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
|
# Migrate data back: wrap interface_type in array
|
||||||
op.execute("""
|
op.execute("""
|
||||||
UPDATE tool_types
|
UPDATE tool_types
|
||||||
SET interfaces = jsonb_build_array(interface_type)
|
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
|
# Drop new columns
|
||||||
op.drop_column('tool_types', 'requires_port')
|
op.drop_column('tool_types', 'requires_port')
|
||||||
|
|||||||
Reference in New Issue
Block a user