From 366ce591cbd6ee92dc21ab5dae6d57b57b6ca92f Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 29 Apr 2024 09:53:56 -0300 Subject: [PATCH] Add default_fields column to the variable table --- .../1f4d6df60295_add_default_fields_column.py | 45 +++++++++++++++++++ src/backend/base/langflow/api/v1/variable.py | 4 +- .../database/models/variable/model.py | 7 ++- 3 files changed, 52 insertions(+), 4 deletions(-) create mode 100644 src/backend/base/langflow/alembic/versions/1f4d6df60295_add_default_fields_column.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 new file mode 100644 index 000000000..7b69058c2 --- /dev/null +++ b/src/backend/base/langflow/alembic/versions/1f4d6df60295_add_default_fields_column.py @@ -0,0 +1,45 @@ +"""Add default_fields column + +Revision ID: 1f4d6df60295 +Revises: 58b28437a398 +Create Date: 2024-04-29 09:49:46.864145 + +""" + +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 = "1f4d6df60295" +down_revision: Union[str, None] = "58b28437a398" +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 + 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: + if "default_fields" not in column_names: + batch_op.add_column(sa.Column("default_fields", sa.JSON(), 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() + # ### 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: + if "default_fields" in column_names: + batch_op.drop_column("default_fields") + + # ### end Alembic commands ### diff --git a/src/backend/base/langflow/api/v1/variable.py b/src/backend/base/langflow/api/v1/variable.py index 123086676..b1cca14c1 100644 --- a/src/backend/base/langflow/api/v1/variable.py +++ b/src/backend/base/langflow/api/v1/variable.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import datetime, timezone from uuid import UUID from fastapi import APIRouter, Depends, HTTPException @@ -85,7 +85,7 @@ def update_variable( variable_data = variable.model_dump(exclude_unset=True) for key, value in variable_data.items(): setattr(db_variable, key, value) - db_variable.updated_at = datetime.utcnow() + db_variable.updated_at = datetime.now(timezone.utc) session.commit() session.refresh(db_variable) return db_variable 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 1fab92b4c..cf2236fe1 100644 --- a/src/backend/base/langflow/services/database/models/variable/model.py +++ b/src/backend/base/langflow/services/database/models/variable/model.py @@ -1,8 +1,8 @@ from datetime import datetime, timezone -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, List, Optional from uuid import UUID, uuid4 -from sqlmodel import Column, DateTime, Field, Relationship, SQLModel, func +from sqlmodel import JSON, Column, DateTime, Field, Relationship, SQLModel, func if TYPE_CHECKING: from langflow.services.database.models.user.model import User @@ -15,6 +15,7 @@ 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") + default_fields: Optional[List[str]] = Field(sa_column=Column(JSON)) type: Optional[str] = Field(None, description="Type of the variable") @@ -35,6 +36,7 @@ class Variable(VariableBase, table=True): sa_column=Column(DateTime(timezone=True), nullable=True), description="Last update time of the variable", ) + default_fields: Optional[List[str]] = Field(sa_column=Column(JSON)) # foreign key to user table user_id: UUID = Field(description="User ID associated with this variable", foreign_key="user.id") user: "User" = Relationship(back_populates="variables") @@ -56,3 +58,4 @@ class VariableUpdate(SQLModel): id: UUID # Include the ID for updating name: Optional[str] = Field(None, description="Name of the variable") value: Optional[str] = Field(None, description="Encrypted value of the variable") + default_fields: Optional[List[str]] = Field(None, description="Default fields for the variable")