diff --git a/src/backend/base/langflow/components/data/sql_executor.py b/src/backend/base/langflow/components/data/sql_executor.py index 582d61cf5..741fe9a5a 100644 --- a/src/backend/base/langflow/components/data/sql_executor.py +++ b/src/backend/base/langflow/components/data/sql_executor.py @@ -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 diff --git a/src/backend/tests/unit/components/data/test_sql_executor.py b/src/backend/tests/unit/components/data/test_sql_executor.py index 14ee99020..b312d85bd 100644 --- a/src/backend/tests/unit/components/data/test_sql_executor.py +++ b/src/backend/tests/unit/components/data/test_sql_executor.py @@ -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"