Skip to content
Open
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -440,6 +440,13 @@ external = ["RUF027"]
"SLF001",
"BLE001", # allow broad-exception catching in tests
]
"src/backend/base/langflow/tests/*" = [
"D1",
"PLR2004",
"S101",
"SLF001",
"BLE001", # allow broad-exception catching in tests
]
"src/lfx/tests/*" = [
"D1",
"PLR2004",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,14 @@ class TransactionBase(SQLModel):
class Config:
arbitrary_types_allowed = True

def __init__(self, **data):
# Filter out the 'code' key from inputs before creating the model
if "inputs" in data and isinstance(data["inputs"], dict) and "code" in data["inputs"]:
# Create a shallow copy and remove the code key
data["inputs"] = data["inputs"].copy()
data["inputs"].pop("code")
super().__init__(**data)

@field_validator("flow_id", mode="before")
@classmethod
def validate_flow_id(cls, value):
Expand All @@ -40,6 +48,12 @@ def serialize_inputs(self, data) -> dict:
Returns:
dict: The serialized input data with applied constraints.
"""
# Filter out the "code" key from inputs if present
if isinstance(data, dict) and "code" in data:
# Create a shallow copy and remove the code key
data = data.copy()
data.pop("code")

return serialize(data, max_length=get_max_text_length(), max_items=get_max_items_length())

@field_serializer("outputs")
Expand Down
3 changes: 3 additions & 0 deletions src/backend/base/langflow/tests/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Tests package for langflow."""

# Made with Bob
3 changes: 3 additions & 0 deletions src/backend/base/langflow/tests/services/database/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Database tests package."""

# Made with Bob
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Database models tests package."""

# Made with Bob
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Transaction models tests package."""

# Made with Bob
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
import uuid
from datetime import datetime, timezone

import pytest
from langflow.services.database.models.transactions.model import TransactionBase


def test_serialize_inputs_excludes_code_key():
"""Test that the code key is excluded from inputs when serializing."""
# Create a TransactionBase object with inputs containing a code key
transaction = TransactionBase(
timestamp=datetime.now(timezone.utc),
vertex_id="test-vertex",
target_id="test-target",
inputs={"param1": "value1", "param2": "value2", "code": "print('Hello, world!')"},
outputs={"result": "success"},
status="completed",
flow_id=uuid.uuid4(),
)

# Get the serialized inputs
serialized_inputs = transaction.serialize_inputs(transaction.inputs)

# Verify that the code key is excluded
assert "code" not in serialized_inputs
assert "param1" in serialized_inputs
assert "param2" in serialized_inputs
assert serialized_inputs["param1"] == "value1"
assert serialized_inputs["param2"] == "value2"


def test_serialize_inputs_handles_none():
"""Test that the serialize_inputs method handles None inputs."""
# Create a TransactionBase object with None inputs
transaction = TransactionBase(
timestamp=datetime.now(timezone.utc),
vertex_id="test-vertex",
target_id="test-target",
inputs=None,
outputs={"result": "success"},
status="completed",
flow_id=uuid.uuid4(),
)

# Get the serialized inputs
serialized_inputs = transaction.serialize_inputs(transaction.inputs)

# Verify that None is returned
assert serialized_inputs is None


def test_serialize_inputs_handles_non_dict():
"""Test that the serialize_inputs method handles non-dict inputs."""
# Create a TransactionBase object with valid inputs
transaction = TransactionBase(
timestamp=datetime.now(timezone.utc),
vertex_id="test-vertex",
target_id="test-target",
inputs={}, # Empty dict is valid
outputs={"result": "success"},
status="completed",
flow_id=uuid.uuid4(),
)

# Call serialize_inputs directly with a non-dict value
serialized_inputs = transaction.serialize_inputs("not a dict")

# Verify that the input is returned as is
assert serialized_inputs == "not a dict"


def test_serialize_inputs_handles_empty_dict():
"""Test that the serialize_inputs method handles empty dict inputs."""
# Create a TransactionBase object with empty dict inputs
transaction = TransactionBase(
timestamp=datetime.now(timezone.utc),
vertex_id="test-vertex",
target_id="test-target",
inputs={},
outputs={"result": "success"},
status="completed",
flow_id=uuid.uuid4(),
)

# Get the serialized inputs
serialized_inputs = transaction.serialize_inputs(transaction.inputs)

# Verify that an empty dict is returned
assert serialized_inputs == {}


@pytest.mark.asyncio
async def test_code_key_not_saved_to_database():
"""Test that the code key is not saved to the database."""
# Create input data with a code key
input_data = {"param1": "value1", "param2": "value2", "code": "print('Hello, world!')"}

# Create a transaction with inputs containing a code key
transaction = TransactionBase(
timestamp=datetime.now(timezone.utc),
vertex_id="test-vertex",
target_id="test-target",
inputs=input_data,
outputs={"result": "success"},
status="completed",
flow_id=uuid.uuid4(),
)

# Verify that the code key is removed during transaction creation
assert transaction.inputs is not None
assert "code" not in transaction.inputs
assert "param1" in transaction.inputs
assert "param2" in transaction.inputs

# Verify that the code key is excluded when serializing
serialized_inputs = transaction.serialize_inputs(transaction.inputs)
assert "code" not in serialized_inputs
assert "param1" in serialized_inputs
assert "param2" in serialized_inputs


# Made with Bob
43 changes: 43 additions & 0 deletions src/backend/tests/unit/test_database.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import json
from datetime import datetime, timezone
from typing import NamedTuple
from uuid import UUID, uuid4

Expand Down Expand Up @@ -778,3 +779,45 @@ async def test_read_folder_with_component_filter(client: AsyncClient, json_flow:
assert len(folder_data["flows"]["items"]) == 1
assert folder_data["flows"]["items"][0]["name"] == component_flow_name
assert folder_data["flows"]["items"][0]["is_component"] == True # noqa: E712


def test_transaction_excludes_code_key(session, active_user):
"""Test that the code key is excluded from transaction inputs when logged to the database."""
from langflow.services.database.models.transactions.model import TransactionTable

# Create a flow to associate with the transaction
flow = Flow(name=str(uuid4()), description="Test flow", data={}, user_id=active_user.id)
session.add(flow)
session.commit()
session.refresh(flow)

# Create input data with a code key
input_data = {"param1": "value1", "param2": "value2", "code": "print('Hello, world!')"}

# Create a transaction with inputs containing a code key
transaction = TransactionTable(
timestamp=datetime.now(timezone.utc),
vertex_id="test-vertex",
target_id="test-target",
inputs=input_data,
outputs={"result": "success"},
status="completed",
flow_id=flow.id,
)

# Verify that the code key is removed during transaction creation
assert transaction.inputs is not None
assert "code" not in transaction.inputs
assert "param1" in transaction.inputs
assert "param2" in transaction.inputs

# Add the transaction to the database
session.add(transaction)
session.commit()
session.refresh(transaction)

# Verify that the code key is not in the saved transaction inputs
assert transaction.inputs is not None
assert "code" not in transaction.inputs
assert "param1" in transaction.inputs
assert "param2" in transaction.inputs
Loading