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.custom.custom_component.component_with_cache import ComponentWithCache
from langflow.io import BoolInput, MessageTextInput, Output from langflow.io import BoolInput, MessageTextInput, Output
from langflow.schema.data import Data
from langflow.schema.dataframe import DataFrame from langflow.schema.dataframe import DataFrame
from langflow.schema.message import Message from langflow.schema.message import Message
from langflow.services.cache.utils import CacheMiss from langflow.services.cache.utils import CacheMiss
@ -43,20 +42,19 @@ class SQLComponent(ComponentWithCache):
inputs = [ inputs = [
MessageTextInput(name="database_url", display_name="Database URL", required=True), MessageTextInput(name="database_url", display_name="Database URL", required=True),
MessageTextInput(name="query", display_name="SQL Query", tool_mode=True, 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( BoolInput(
name="add_error", name="add_error",
display_name="Add Error", display_name="Add Error",
value=False, value=False,
tool_mode=True, tool_mode=True,
info="If True, the error will be added to the result", info="If True, the error will be added to the result",
advanced=True,
), ),
] ]
outputs = [ outputs = [
Output(display_name="Message", name="text", method="build_component"), Output(display_name="Query Results", name="sql_query_results", method="sql_query_results"),
Output(display_name="Data", name="data", method="build_data"),
Output(display_name="DataFrame", name="dataframe", method="build_dataframe"),
] ]
def build_component( def build_component(
@ -92,14 +90,8 @@ class SQLComponent(ComponentWithCache):
self.log(msg) self.log(msg)
raise ValueError(msg) from e raise ValueError(msg) from e
def build_dataframe(self) -> DataFrame: def sql_query_results(self) -> DataFrame:
result = self.__execute_query() result = self.__execute_query()
df_result = DataFrame(result) df_result = DataFrame(result)
self.status = df_result self.status = df_result
return 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 import pytest
from langflow.components.data.sql_executor import SQLComponent 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 from tests.base import ComponentTestBaseWithoutClient
@ -89,11 +89,11 @@ class TestSQLComponent(ComponentTestBaseWithoutClient):
assert "Error:" in result.text assert "Error:" in result.text
assert "Query: SELECT * FROM non_existent_table" 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.""" """Test building a DataFrame from a SQL query."""
component = component_class(**default_kwargs) component = component_class(**default_kwargs)
result = component.build_dataframe() result = component.sql_query_results()
assert isinstance(result, DataFrame) assert isinstance(result, DataFrame)
assert len(result) == 1 assert len(result) == 1
@ -101,15 +101,3 @@ class TestSQLComponent(ComponentTestBaseWithoutClient):
assert "name" in result.columns assert "name" in result.columns
assert result.iloc[0]["id"] == 1 assert result.iloc[0]["id"] == 1
assert result.iloc[0]["name"] == "name_test" 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"