From 369126a125cfaab3cd425c7fc3073a2662832050 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Thu, 5 Dec 2024 16:11:17 -0300 Subject: [PATCH] refactor: use Inspector for column checks in locked state migration (#5097) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🐛 (add_persistent_locked_state.py): fix migration script to correctly check for existing 'locked' column before adding or dropping it --- ...e3162c1804e6_add_persistent_locked_state.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/backend/base/langflow/alembic/versions/e3162c1804e6_add_persistent_locked_state.py b/src/backend/base/langflow/alembic/versions/e3162c1804e6_add_persistent_locked_state.py index 80de87a35..9bfa34eb9 100644 --- a/src/backend/base/langflow/alembic/versions/e3162c1804e6_add_persistent_locked_state.py +++ b/src/backend/base/langflow/alembic/versions/e3162c1804e6_add_persistent_locked_state.py @@ -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 ###