From 050230cad1a6e9f67de8291e7241f51198aeb567 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 11 Apr 2025 17:24:38 -0300 Subject: [PATCH] fix: Enhance foreign key constraint management and naming conventions in migrations (#7569) * fix: ensure foreign key constraints are only dropped if they exist in upgrade script * Updated the upgrade function to check for the existence of foreign key constraint names before attempting to drop them in the message, transaction, and vertex_build tables, enhancing robustness and preventing potential errors during migration. * refactor: streamline foreign key constraint removal in upgrade script * Simplified the upgrade function by directly dropping foreign key constraints for the message, transaction, and vertex_build tables without checking for their existence, enhancing code clarity and maintainability. * feat: introduce naming convention for database constraints in Alembic environment * Added a naming convention dictionary to standardize the naming of indexes, unique constraints, check constraints, foreign keys, and primary keys in the Alembic migration environment, enhancing consistency and clarity in database schema management. * refactor: enhance foreign key constraint management in upgrade and downgrade scripts * Updated the upgrade function to recreate the message, transaction, and vertex_build tables without foreign key constraints, preserving data integrity during migration. * Improved the downgrade function to restore these tables with the appropriate foreign key constraints, ensuring consistency in the database schema. * Introduced a naming convention for database constraints to standardize naming across migrations. * refactor: update table schema in Alembic migration for improved data handling * Modified the schema for the message, transaction, and vertex_build tables to enhance data integrity by changing column types and adding new fields. * Updated data insertion queries to use quoted identifiers, preventing potential issues with SQL keywords. * Ensured that the upgrade and downgrade functions maintain consistency in the database schema during migrations. * refactor: enhance data insertion logic in Alembic migration scripts * Updated data insertion queries for the message, transaction, and vertex_build tables to explicitly list columns and filter out rows with NULL IDs, improving data integrity during migrations. * Utilized window functions to ensure uniqueness of build_id across different database systems, enhancing compatibility and robustness. * Maintained consistency in upgrade and downgrade functions to ensure seamless schema transitions. * Revert "Revert "fix: published flows now can add rows to the database" (#7571)" This reverts commit a8d2fccd1c782e73e8594d08ea1bbfae7959a804. --- src/backend/base/langflow/alembic/env.py | 9 +- ...0a6fa3_remove_fk_constraint_in_message_.py | 271 ++++++++++++++++++ .../services/database/models/__init__.py | 3 +- .../services/database/models/flow/model.py | 6 - .../services/database/models/message/model.py | 6 +- .../database/models/transactions/model.py | 9 +- .../database/models/vertex_builds/model.py | 9 +- 7 files changed, 287 insertions(+), 26 deletions(-) create mode 100644 src/backend/base/langflow/alembic/versions/1b8b740a6fa3_remove_fk_constraint_in_message_.py diff --git a/src/backend/base/langflow/alembic/env.py b/src/backend/base/langflow/alembic/env.py index 31b88d875..8419403ab 100644 --- a/src/backend/base/langflow/alembic/env.py +++ b/src/backend/base/langflow/alembic/env.py @@ -18,12 +18,19 @@ config = context.config if config.config_file_name is not None: fileConfig(config.config_file_name) +NAMING_CONVENTION = { + "ix": "ix_%(column_0_label)s", + "uq": "uq_%(table_name)s_%(column_0_name)s", + "ck": "ck_%(table_name)s_%(constraint_name)s", + "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", + "pk": "pk_%(table_name)s", +} # add your model's MetaData object here # for 'autogenerate' support # from myapp import mymodel # target_metadata = mymodel.Base.metadata target_metadata = SQLModel.metadata - +target_metadata.naming_convention = NAMING_CONVENTION # other values from the config, defined by the needs of env.py, # can be acquired: # my_important_option = config.get_main_option("my_important_option") diff --git a/src/backend/base/langflow/alembic/versions/1b8b740a6fa3_remove_fk_constraint_in_message_.py b/src/backend/base/langflow/alembic/versions/1b8b740a6fa3_remove_fk_constraint_in_message_.py new file mode 100644 index 000000000..b385f9bd9 --- /dev/null +++ b/src/backend/base/langflow/alembic/versions/1b8b740a6fa3_remove_fk_constraint_in_message_.py @@ -0,0 +1,271 @@ +"""remove fk constraint in message transaction and vertex build + + +Revision ID: 1b8b740a6fa3 +Revises: f3b2d1f1002d +Create Date: 2025-04-10 10:17:32.493181 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel +from sqlalchemy.engine.reflection import Inspector +from langflow.utils import migration + + +# revision identifiers, used by Alembic. +revision: str = '1b8b740a6fa3' +down_revision: Union[str, None] = 'f3b2d1f1002d' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +NAMING_CONVENTION = { + "ix": "ix_%(column_0_label)s", + "uq": "uq_%(table_name)s_%(column_0_name)s", + "ck": "ck_%(table_name)s_%(constraint_name)s", + "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s", + "pk": "pk_%(table_name)s", +} + +def upgrade() -> None: + conn = op.get_bind() + + # For SQLite, we need to recreate the tables without the constraints + # This approach preserves all data while removing the constraints + + # 1. Handle vertex_build table + if migration.table_exists("vertex_build", conn): + # Create a temporary table without the constraint + temp_table_name = "temp_vertex_build" + + # Create temp table with same schema but no FK constraint + op.create_table( + temp_table_name, + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("data", sa.JSON(), nullable=True), + sa.Column("artifacts", sa.JSON(), nullable=True), + sa.Column("params", sa.Text(), nullable=True), + sa.Column("build_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("valid", sa.BOOLEAN(), nullable=False), + sa.PrimaryKeyConstraint("build_id", name="pk_vertex_build"), + ) + + # Copy data - use a window function to ensure build_id uniqueness across SQLite, PostgreSQL and MySQL + # Filter out rows where the original 'id' (vertex id) is NULL, as the new table requires it. + op.execute(f''' + INSERT INTO "{temp_table_name}" (timestamp, id, data, artifacts, params, build_id, flow_id, valid) + SELECT timestamp, id, data, artifacts, params, build_id, flow_id, valid + FROM ( + SELECT timestamp, id, data, artifacts, params, build_id, flow_id, valid, + ROW_NUMBER() OVER (PARTITION BY build_id ORDER BY timestamp) as rn + FROM "vertex_build" + WHERE id IS NOT NULL -- Ensure vertex id is not NULL + ) sub + WHERE rn = 1 + ''') + + # Drop original table and rename temp table + op.drop_table("vertex_build") + op.rename_table(temp_table_name, "vertex_build") + + # 2. Handle transaction table + if migration.table_exists("transaction", conn): + # Create a temporary table without the constraint + temp_table_name = "temp_transaction" + + # Create temp table with same schema but no FK constraint + op.create_table( + temp_table_name, + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("vertex_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("target_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("inputs", sa.JSON(), nullable=True), + sa.Column("outputs", sa.JSON(), nullable=True), + sa.Column("status", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("error", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.PrimaryKeyConstraint("id", name="pk_transaction"), + ) + + # Copy data - explicitly list columns and filter out rows where id is NULL + op.execute(f''' + INSERT INTO "{temp_table_name}" (timestamp, vertex_id, target_id, inputs, outputs, status, id, flow_id, error) + SELECT timestamp, vertex_id, target_id, inputs, outputs, status, id, flow_id, error + FROM "transaction" + WHERE id IS NOT NULL + ''') + + # Drop original table and rename temp table + op.drop_table("transaction") + op.rename_table(temp_table_name, "transaction") + + # 3. Handle message table + if migration.table_exists("message", conn): + # Create a temporary table without the constraint + temp_table_name = "temp_message" + + # Create temp table with same schema but no FK constraint + op.create_table( + temp_table_name, + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("sender", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("sender_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("session_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("text", sa.Text(), nullable=True), + sa.Column("id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=True), + sa.Column("files", sa.JSON(), nullable=True), + sa.Column("error", sa.Boolean(), nullable=False, server_default=sa.false()), + sa.Column("edit", sa.Boolean(), nullable=False, server_default=sa.false()), + sa.Column("properties", sa.JSON(), nullable=True), + sa.Column("category", sa.Text(), nullable=True), + sa.Column("content_blocks", sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint("id", name="pk_message"), + ) + + # Copy data - explicitly list columns and filter out rows where id is NULL + op.execute(f''' + INSERT INTO "{temp_table_name}" (timestamp, sender, sender_name, session_id, text, id, flow_id, files, error, edit, properties, category, content_blocks) + SELECT timestamp, sender, sender_name, session_id, text, id, flow_id, files, error, edit, properties, category, content_blocks + FROM "message" + WHERE id IS NOT NULL + ''') + + # Drop original table and rename temp table + op.drop_table("message") + op.rename_table(temp_table_name, "message") + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + conn = op.get_bind() + + # Recreate tables with constraints + # 1. Handle vertex_build table + if migration.table_exists("vertex_build", conn): + # Create a temporary table with the constraint + temp_table_name = "temp_vertex_build" + + # Create temp table with same schema including FK constraint + # Note: Original 'id' column was nullable=True here, which might be inconsistent with the model. + # Keeping it as nullable=True to match the previous state of this downgrade function. + op.create_table( + temp_table_name, + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("data", sa.JSON(), nullable=True), + sa.Column("artifacts", sa.JSON(), nullable=True), + sa.Column("params", sa.Text(), nullable=True), + sa.Column("build_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("valid", sa.BOOLEAN(), nullable=False), + sa.ForeignKeyConstraint( + ["flow_id"], + ["flow.id"], + name="fk_vertex_build_flow_id_flow", + ), + sa.PrimaryKeyConstraint("build_id", name="pk_vertex_build"), + ) + + # Copy data - use a window function to ensure build_id uniqueness. + # Filter out rows where build_id is NULL (PK constraint) + # No need to filter by 'id' here as the target column allows NULLs. + op.execute(f''' + INSERT INTO "{temp_table_name}" (timestamp, id, data, artifacts, params, build_id, flow_id, valid) + SELECT timestamp, id, data, artifacts, params, build_id, flow_id, valid + FROM ( + SELECT timestamp, id, data, artifacts, params, build_id, flow_id, valid, + ROW_NUMBER() OVER (PARTITION BY build_id ORDER BY timestamp) as rn + FROM "vertex_build" + WHERE build_id IS NOT NULL -- Ensure primary key is not NULL + ) sub + WHERE rn = 1 + ''') + + # Drop original table and rename temp table + op.drop_table("vertex_build") + op.rename_table(temp_table_name, "vertex_build") + + # 2. Handle transaction table + if migration.table_exists("transaction", conn): + # Create a temporary table with the constraint + temp_table_name = "temp_transaction" + + # Create temp table with same schema including FK constraint + op.create_table( + temp_table_name, + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("vertex_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("target_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("inputs", sa.JSON(), nullable=True), + sa.Column("outputs", sa.JSON(), nullable=True), + sa.Column("status", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("error", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.ForeignKeyConstraint( + ["flow_id"], + ["flow.id"], + name="fk_transaction_flow_id_flow", + ), + sa.PrimaryKeyConstraint("id", name="pk_transaction"), + ) + + # Copy data - explicitly list columns and filter out rows where id is NULL + op.execute(f''' + INSERT INTO "{temp_table_name}" (timestamp, vertex_id, target_id, inputs, outputs, status, id, flow_id, error) + SELECT timestamp, vertex_id, target_id, inputs, outputs, status, id, flow_id, error + FROM "transaction" + WHERE id IS NOT NULL + ''') + + # Drop original table and rename temp table + op.drop_table("transaction") + op.rename_table(temp_table_name, "transaction") + + # 3. Handle message table + if migration.table_exists("message", conn): + # Create a temporary table with the constraint + temp_table_name = "temp_message" + + # Create temp table with same schema including FK constraint + op.create_table( + temp_table_name, + sa.Column("timestamp", sa.DateTime(), nullable=False), + sa.Column("sender", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("sender_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("session_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("text", sa.Text(), nullable=True), + sa.Column("id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=False), + sa.Column("flow_id", sqlmodel.sql.sqltypes.types.Uuid(), nullable=True), + sa.Column("files", sa.JSON(), nullable=True), + sa.Column("error", sa.Boolean(), nullable=False, server_default=sa.false()), + sa.Column("edit", sa.Boolean(), nullable=False, server_default=sa.false()), + sa.Column("properties", sa.JSON(), nullable=True), + sa.Column("category", sa.Text(), nullable=True), + sa.Column("content_blocks", sa.JSON(), nullable=True), + sa.ForeignKeyConstraint( + ["flow_id"], + ["flow.id"], + name="fk_message_flow_id_flow", + ), + sa.PrimaryKeyConstraint("id", name="pk_message"), + ) + + # Copy data - explicitly list columns and filter out rows where id is NULL + op.execute(f''' + INSERT INTO "{temp_table_name}" (timestamp, sender, sender_name, session_id, text, id, flow_id, files, error, edit, properties, category, content_blocks) + SELECT timestamp, sender, sender_name, session_id, text, id, flow_id, files, error, edit, properties, category, content_blocks + FROM "message" + WHERE id IS NOT NULL + ''') + + # Drop original table and rename temp table + op.drop_table("message") + op.rename_table(temp_table_name, "message") + # ### end Alembic commands ### diff --git a/src/backend/base/langflow/services/database/models/__init__.py b/src/backend/base/langflow/services/database/models/__init__.py index 4419e7f11..d8247c754 100644 --- a/src/backend/base/langflow/services/database/models/__init__.py +++ b/src/backend/base/langflow/services/database/models/__init__.py @@ -1,4 +1,5 @@ from .api_key import ApiKey +from .file import File from .flow import Flow from .folder import Folder from .message import MessageTable @@ -6,4 +7,4 @@ from .transactions import TransactionTable from .user import User from .variable import Variable -__all__ = ["ApiKey", "Flow", "Folder", "MessageTable", "TransactionTable", "User", "Variable"] +__all__ = ["ApiKey", "File", "Flow", "Folder", "MessageTable", "TransactionTable", "User", "Variable"] diff --git a/src/backend/base/langflow/services/database/models/flow/model.py b/src/backend/base/langflow/services/database/models/flow/model.py index a706c7ad9..cca40ecda 100644 --- a/src/backend/base/langflow/services/database/models/flow/model.py +++ b/src/backend/base/langflow/services/database/models/flow/model.py @@ -23,11 +23,8 @@ from sqlmodel import JSON, Column, Field, Relationship, SQLModel from langflow.schema import Data if TYPE_CHECKING: - from langflow.services.database.models import TransactionTable from langflow.services.database.models.folder import Folder - from langflow.services.database.models.message import MessageTable from langflow.services.database.models.user import User - from langflow.services.database.models.vertex_builds.model import VertexBuildTable HEX_COLOR_LENGTH = 7 @@ -190,9 +187,6 @@ class Flow(FlowBase, table=True): # type: ignore[call-arg] folder_id: UUID | None = Field(default=None, foreign_key="folder.id", nullable=True, index=True) fs_path: str | None = Field(default=None, nullable=True) folder: Optional["Folder"] = Relationship(back_populates="flows") - messages: list["MessageTable"] = Relationship(back_populates="flow") - transactions: list["TransactionTable"] = Relationship(back_populates="flow") - vertex_builds: list["VertexBuildTable"] = Relationship(back_populates="flow") def to_data(self): serialized = self.model_dump() diff --git a/src/backend/base/langflow/services/database/models/message/model.py b/src/backend/base/langflow/services/database/models/message/model.py index df21aa8e4..f0b3fdb77 100644 --- a/src/backend/base/langflow/services/database/models/message/model.py +++ b/src/backend/base/langflow/services/database/models/message/model.py @@ -5,7 +5,7 @@ from uuid import UUID, uuid4 from pydantic import field_serializer, field_validator from sqlalchemy import Text -from sqlmodel import JSON, Column, Field, Relationship, SQLModel +from sqlmodel import JSON, Column, Field, SQLModel from langflow.schema.content_block import ContentBlock from langflow.schema.properties import Properties @@ -13,7 +13,6 @@ from langflow.schema.validators import str_to_timestamp_validator if TYPE_CHECKING: from langflow.schema.message import Message - from langflow.services.database.models.flow.model import Flow class MessageBase(SQLModel): @@ -113,8 +112,7 @@ class MessageBase(SQLModel): class MessageTable(MessageBase, table=True): # type: ignore[call-arg] __tablename__ = "message" id: UUID = Field(default_factory=uuid4, primary_key=True) - flow_id: UUID | None = Field(default=None, foreign_key="flow.id") - flow: "Flow" = Relationship(back_populates="messages") + flow_id: UUID | None = Field(default=None) files: list[str] = Field(sa_column=Column(JSON)) properties: Properties = Field(default_factory=lambda: Properties().model_dump(), sa_column=Column(JSON)) # type: ignore[assignment] category: str = Field(sa_column=Column(Text)) diff --git a/src/backend/base/langflow/services/database/models/transactions/model.py b/src/backend/base/langflow/services/database/models/transactions/model.py index eb43ae5b4..627bf7101 100644 --- a/src/backend/base/langflow/services/database/models/transactions/model.py +++ b/src/backend/base/langflow/services/database/models/transactions/model.py @@ -1,16 +1,12 @@ from datetime import datetime, timezone -from typing import TYPE_CHECKING from uuid import UUID, uuid4 from pydantic import field_serializer, field_validator -from sqlmodel import JSON, Column, Field, Relationship, SQLModel +from sqlmodel import JSON, Column, Field, SQLModel from langflow.serialization.constants import MAX_ITEMS_LENGTH, MAX_TEXT_LENGTH from langflow.serialization.serialization import serialize -if TYPE_CHECKING: - from langflow.services.database.models.flow.model import Flow - class TransactionBase(SQLModel): timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) @@ -20,7 +16,7 @@ class TransactionBase(SQLModel): outputs: dict | None = Field(default=None, sa_column=Column(JSON)) status: str = Field(nullable=False) error: str | None = Field(default=None) - flow_id: UUID = Field(foreign_key="flow.id") + flow_id: UUID = Field() # Needed for Column(JSON) class Config: @@ -47,7 +43,6 @@ class TransactionBase(SQLModel): class TransactionTable(TransactionBase, table=True): # type: ignore[call-arg] __tablename__ = "transaction" id: UUID | None = Field(default_factory=uuid4, primary_key=True) - flow: "Flow" = Relationship(back_populates="transactions") class TransactionReadResponse(TransactionBase): diff --git a/src/backend/base/langflow/services/database/models/vertex_builds/model.py b/src/backend/base/langflow/services/database/models/vertex_builds/model.py index 38e5dba0b..411438099 100644 --- a/src/backend/base/langflow/services/database/models/vertex_builds/model.py +++ b/src/backend/base/langflow/services/database/models/vertex_builds/model.py @@ -1,17 +1,13 @@ from datetime import datetime, timezone -from typing import TYPE_CHECKING from uuid import UUID, uuid4 from pydantic import BaseModel, field_serializer, field_validator from sqlalchemy import Text -from sqlmodel import JSON, Column, Field, Relationship, SQLModel +from sqlmodel import JSON, Column, Field, SQLModel from langflow.serialization.constants import MAX_ITEMS_LENGTH, MAX_TEXT_LENGTH from langflow.serialization.serialization import serialize -if TYPE_CHECKING: - from langflow.services.database.models.flow.model import Flow - class VertexBuildBase(SQLModel): timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) @@ -20,7 +16,7 @@ class VertexBuildBase(SQLModel): artifacts: dict | None = Field(default=None, sa_column=Column(JSON)) params: str | None = Field(default=None, sa_column=Column(Text, nullable=True)) valid: bool = Field(nullable=False) - flow_id: UUID = Field(foreign_key="flow.id") + flow_id: UUID = Field() # Needed for Column(JSON) class Config: @@ -58,7 +54,6 @@ class VertexBuildBase(SQLModel): class VertexBuildTable(VertexBuildBase, table=True): # type: ignore[call-arg] __tablename__ = "vertex_build" build_id: UUID | None = Field(default_factory=uuid4, primary_key=True) - flow: "Flow" = Relationship(back_populates="vertex_builds") class VertexBuildMapModel(BaseModel):