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:
Cristhian Zanforlin Lousa 2024-12-05 16:11:17 -03:00 committed by GitHub
commit 369126a125
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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 ###