jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: Enhance migration script with column existence check and improved error handling (#4680)

Browse source 
* Enhance migration script with column existence check and improved error handling

* Refactor migration script to use `column_exists` with explicit parameters
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-11-18 17:21:58 -03:00 • committed by GitHub
commit 160409d19c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 24 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

34
src/backend/base/langflow/alembic/versions/1ef9c4f3765d_.py
View file

@ -12,6 +12,9 @@ from typing import Sequence, Union
import sqlalchemy as sa import sqlalchemy as sa
import sqlmodel import sqlmodel
from alembic import op from alembic import op
from sqlalchemy.engine.reflection import Inspector
from langflow.utils import migration
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = "1ef9c4f3765d" revision: str = "1ef9c4f3765d"
@ -21,20 +24,31 @@ depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None: def upgrade() -> None:
conn = op.get_bind()
# ### commands auto generated by Alembic - please adjust! ### # ### commands auto generated by Alembic - please adjust! ###
try: inspector = Inspector.from_engine(conn) # type: ignore
with op.batch_alter_table("apikey", schema=None) as batch_op: # ### commands auto generated by Alembic - please adjust! ###
batch_op.alter_column("name", existing_type=sqlmodel.sql.sqltypes.AutoString(), nullable=True)
except Exception: with op.batch_alter_table("apikey", schema=None) as batch_op:
pass if migration.column_exists(table_name="apikey", column_name="name", conn=conn):
api_key_columns = inspector.get_columns("apikey")
name_column = next((column for column in api_key_columns if column["name"] == "name"), None)
if name_column is not None and isinstance(name_column["type"], sa.VARCHAR) and not name_column["nullable"]:
batch_op.alter_column("name", existing_type=sa.VARCHAR(), nullable=True)
# ### end Alembic commands ### # ### end Alembic commands ###
def downgrade() -> None: def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ### # ### commands auto generated by Alembic - please adjust! ###
try: conn = op.get_bind()
with op.batch_alter_table("apikey", schema=None) as batch_op: inspector = Inspector.from_engine(conn) # type: ignore
batch_op.alter_column("name", existing_type=sa.VARCHAR(), nullable=False)
except Exception: with op.batch_alter_table("apikey", schema=None) as batch_op:
pass if migration.column_exists(table_name="apikey", column_name="name", conn=conn):
api_key_columns = inspector.get_columns("apikey")
name_column = next((column for column in api_key_columns if column["name"] == "name"), None)
if name_column is not None and isinstance(name_column["type"], sa.VARCHAR) and name_column["nullable"]:
batch_op.alter_column("name", existing_type=sa.VARCHAR(), nullable=False)
# ### end Alembic commands ### # ### end Alembic commands ###