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 typing import Sequence, Union
from alembic import op
import sqlalchemy as sa import sqlalchemy as sa
import sqlmodel import sqlmodel
from alembic import op
from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.reflection import Inspector
from langflow.utils import migration
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = 'e3162c1804e6' revision: str = 'e3162c1804e6'
@ -23,13 +21,23 @@ depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None: def upgrade() -> None:
conn = op.get_bind() 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: 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)) batch_op.add_column(sa.Column('locked', sa.Boolean(), nullable=True))
# ### end Alembic commands ###
def downgrade() -> None: def downgrade() -> None:
conn = op.get_bind() 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: 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') batch_op.drop_column('locked')
# ### end Alembic commands ###