refactor: use Inspector for column checks in locked state migration (#5097)
🐛 (add_persistent_locked_state.py): fix migration script to correctly check for existing 'locked' column before adding or dropping it
This commit is contained in:
parent
360fad4287
commit
369126a125
1 changed files with 13 additions and 5 deletions
|
|
@ -7,12 +7,10 @@ Create Date: 2024-11-07 14:50:35.201760
|
|||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
from sqlalchemy.engine.reflection import Inspector
|
||||
from langflow.utils import migration
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'e3162c1804e6'
|
||||
|
|
@ -23,13 +21,23 @@ depends_on: Union[str, Sequence[str], None] = None
|
|||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = Inspector.from_engine(conn) # type: ignore
|
||||
table_names = inspector.get_table_names() # noqa
|
||||
column_names = [column["name"] for column in inspector.get_columns("flow")]
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
with op.batch_alter_table('flow', schema=None) as batch_op:
|
||||
if not migration.column_exists(table_name='flow', column_name='locked', conn=conn):
|
||||
if "locked" not in column_names:
|
||||
batch_op.add_column(sa.Column('locked', sa.Boolean(), nullable=True))
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = Inspector.from_engine(conn) # type: ignore
|
||||
table_names = inspector.get_table_names() # noqa
|
||||
column_names = [column["name"] for column in inspector.get_columns("flow")]
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
with op.batch_alter_table('flow', schema=None) as batch_op:
|
||||
if migration.column_exists(table_name='flow', column_name='locked', conn=conn):
|
||||
if "locked" in column_names:
|
||||
batch_op.drop_column('locked')
|
||||
# ### end Alembic commands ###
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue