feat: improve SQL component UX and Tool mode function names (#7988)

* update component

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Eric Hare <ericrhare@gmail.com>
This commit is contained in:
Edwin Jose 2025-05-16 11:51:06 -04:00 committed by GitHub
commit 777200fab3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 27 deletions

View file

@ -5,7 +5,6 @@ from sqlalchemy.exc import SQLAlchemyError
from langflow.custom.custom_component.component_with_cache import ComponentWithCache
from langflow.io import BoolInput, MessageTextInput, Output
from langflow.schema.data import Data
from langflow.schema.dataframe import DataFrame
from langflow.schema.message import Message
from langflow.services.cache.utils import CacheMiss
@ -43,20 +42,19 @@ class SQLComponent(ComponentWithCache):
inputs = [
MessageTextInput(name="database_url", display_name="Database URL", required=True),
MessageTextInput(name="query", display_name="SQL Query", tool_mode=True, required=True),
BoolInput(name="include_columns", display_name="Include Columns", value=True, tool_mode=True),
BoolInput(name="include_columns", display_name="Include Columns", value=True, tool_mode=True, advanced=True),
BoolInput(
name="add_error",
display_name="Add Error",
value=False,
tool_mode=True,
info="If True, the error will be added to the result",
advanced=True,
),
]
outputs = [
Output(display_name="Message", name="text", method="build_component"),
Output(display_name="Data", name="data", method="build_data"),
Output(display_name="DataFrame", name="dataframe", method="build_dataframe"),
Output(display_name="Query Results", name="sql_query_results", method="sql_query_results"),
]
def build_component(
@ -92,14 +90,8 @@ class SQLComponent(ComponentWithCache):
self.log(msg)
raise ValueError(msg) from e
def build_dataframe(self) -> DataFrame:
def sql_query_results(self) -> DataFrame:
result = self.__execute_query()
df_result = DataFrame(result)
self.status = df_result
return df_result
def build_data(self) -> Data:
result = self.__execute_query()
data_result = Data(data={"result": result})
self.status = data_result
return data_result

View file

@ -3,7 +3,7 @@ from pathlib import Path
import pytest
from langflow.components.data.sql_executor import SQLComponent
from langflow.schema import Data, DataFrame, Message
from langflow.schema import DataFrame, Message
from tests.base import ComponentTestBaseWithoutClient
@ -89,11 +89,11 @@ class TestSQLComponent(ComponentTestBaseWithoutClient):
assert "Error:" in result.text
assert "Query: SELECT * FROM non_existent_table" in result.text
def test_build_dataframe(self, component_class: type[SQLComponent], default_kwargs):
def test_sql_query_results(self, component_class: type[SQLComponent], default_kwargs):
"""Test building a DataFrame from a SQL query."""
component = component_class(**default_kwargs)
result = component.build_dataframe()
result = component.sql_query_results()
assert isinstance(result, DataFrame)
assert len(result) == 1
@ -101,15 +101,3 @@ class TestSQLComponent(ComponentTestBaseWithoutClient):
assert "name" in result.columns
assert result.iloc[0]["id"] == 1
assert result.iloc[0]["name"] == "name_test"
def test_build_data(self, component_class: type[SQLComponent], default_kwargs):
"""Test building a Data object from a SQL query."""
component = component_class(**default_kwargs)
result = component.build_data()
assert isinstance(result, Data)
assert "result" in result.data
assert len(result.data["result"]) == 1
assert result.data["result"][0]["id"] == 1
assert result.data["result"][0]["name"] == "name_test"