From 90f1efb18b8f7b2ac88d79f78aab083c7c347e6d Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 30 Apr 2024 14:34:54 -0300 Subject: [PATCH] Refactor Alembic migration scripts and Variable model --- .../1f4d6df60295_add_default_fields_column.py | 2 - ...d85f_set_name_and_value_to_not_nullable.py | 52 +++++++++++++++++++ .../database/models/variable/model.py | 4 +- 3 files changed, 54 insertions(+), 4 deletions(-) create mode 100644 src/backend/base/langflow/alembic/versions/c153816fd85f_set_name_and_value_to_not_nullable.py diff --git a/src/backend/base/langflow/alembic/versions/1f4d6df60295_add_default_fields_column.py b/src/backend/base/langflow/alembic/versions/1f4d6df60295_add_default_fields_column.py index 7b69058c2..7cf2fd083 100644 --- a/src/backend/base/langflow/alembic/versions/1f4d6df60295_add_default_fields_column.py +++ b/src/backend/base/langflow/alembic/versions/1f4d6df60295_add_default_fields_column.py @@ -22,7 +22,6 @@ 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() # ### commands auto generated by Alembic - please adjust! ### column_names = [column["name"] for column in inspector.get_columns("variable")] with op.batch_alter_table("variable", schema=None) as batch_op: @@ -35,7 +34,6 @@ def upgrade() -> None: def downgrade() -> None: conn = op.get_bind() inspector = Inspector.from_engine(conn) # type: ignore - table_names = inspector.get_table_names() # ### commands auto generated by Alembic - please adjust! ### column_names = [column["name"] for column in inspector.get_columns("variable")] with op.batch_alter_table("variable", schema=None) as batch_op: diff --git a/src/backend/base/langflow/alembic/versions/c153816fd85f_set_name_and_value_to_not_nullable.py b/src/backend/base/langflow/alembic/versions/c153816fd85f_set_name_and_value_to_not_nullable.py new file mode 100644 index 000000000..26cf872ab --- /dev/null +++ b/src/backend/base/langflow/alembic/versions/c153816fd85f_set_name_and_value_to_not_nullable.py @@ -0,0 +1,52 @@ +"""Set name and value to not nullable + +Revision ID: c153816fd85f +Revises: 1f4d6df60295 +Create Date: 2024-04-30 14:31:23.898995 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.engine.reflection import Inspector + +# revision identifiers, used by Alembic. +revision: str = "c153816fd85f" +down_revision: Union[str, None] = "1f4d6df60295" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + conn = op.get_bind() + inspector = Inspector.from_engine(conn) # type: ignore + # ### commands auto generated by Alembic - please adjust! ### + columns = inspector.get_columns("variable") + with op.batch_alter_table("variable", schema=None) as batch_op: + name_column = [column for column in columns if column["name"] == "name"][0] + if name_column and name_column[0]["nullable"]: + batch_op.alter_column("name", existing_type=sa.VARCHAR(), nullable=False) + value_column = [column for column in columns if column["name"] == "value"][0] + if value_column and value_column[0]["nullable"]: + batch_op.alter_column("value", existing_type=sa.VARCHAR(), nullable=False) + + +# ### end Alembic commands ### + + +def downgrade() -> None: + conn = op.get_bind() + inspector = Inspector.from_engine(conn) # type: ignore + columns = inspector.get_columns("variable") + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("variable", schema=None) as batch_op: + name_column = [column for column in columns if column["name"] == "name"][0] + if name_column and not name_column[0]["nullable"]: + batch_op.alter_column("name", existing_type=sa.VARCHAR(), nullable=True) + value_column = [column for column in columns if column["name"] == "value"][0] + if value_column and not value_column[0]["nullable"]: + batch_op.alter_column("name", existing_type=sa.VARCHAR(), nullable=True) + + # ### end Alembic commands ### diff --git a/src/backend/base/langflow/services/database/models/variable/model.py b/src/backend/base/langflow/services/database/models/variable/model.py index 5c56f5b6d..bf61c25dc 100644 --- a/src/backend/base/langflow/services/database/models/variable/model.py +++ b/src/backend/base/langflow/services/database/models/variable/model.py @@ -13,8 +13,8 @@ def utc_now(): class VariableBase(SQLModel): - name: Optional[str] = Field(None, description="Name of the variable") - value: Optional[str] = Field(None, description="Encrypted value of the variable") + name: str = Field(description="Name of the variable") + value: str = Field(description="Encrypted value of the variable") default_fields: Optional[List[str]] = Field(sa_column=Column(JSON)) type: Optional[str] = Field(None, description="Type of the variable")