diff --git a/src/backend/tests/unit/components/data_source/test_sql_executor.py b/src/backend/tests/unit/components/data_source/test_sql_executor.py index 6d8ff886b7..4696e90384 100644 --- a/src/backend/tests/unit/components/data_source/test_sql_executor.py +++ b/src/backend/tests/unit/components/data_source/test_sql_executor.py @@ -1,13 +1,26 @@ import sqlite3 from pathlib import Path +from unittest.mock import Mock, patch import pytest -from lfx.components.data_source.sql_executor import SQLComponent +from lfx.components.data_source.sql_executor import SQL_DATABASE_ENGINE_ARGS, SQLComponent from lfx.schema import DataFrame, Message +from lfx.services.cache.utils import CacheMiss from tests.base import ComponentTestBaseWithoutClient +class FakeSharedComponentCache: + def __init__(self, values=None): + self.values = values or {} + + def get(self, key): + return self.values.get(key, CacheMiss()) + + def set(self, key, value): + self.values[key] = value + + class TestSQLComponent(ComponentTestBaseWithoutClient): @pytest.fixture def test_db(self): @@ -101,3 +114,33 @@ class TestSQLComponent(ComponentTestBaseWithoutClient): assert "name" in result.columns assert result.iloc[0]["id"] == 1 assert result.iloc[0]["name"] == "name_test" + + def test_maybe_create_db_reuses_cached_database(self, component_class: type[SQLComponent], default_kwargs): + """Test cached database reuse without creating a new engine.""" + cached_db = Mock() + cache = FakeSharedComponentCache({default_kwargs["database_url"]: cached_db}) + component = component_class(**default_kwargs) + component._shared_component_cache = cache + + with patch("lfx.components.data_source.sql_executor.SQLDatabase.from_uri") as mock_from_uri: + component.maybe_create_db() + + assert component.db is cached_db + mock_from_uri.assert_not_called() + + def test_maybe_create_db_enables_pool_pre_ping(self, component_class: type[SQLComponent], default_kwargs): + """Test new cached databases validate stale pooled connections on checkout.""" + created_db = Mock() + cache = FakeSharedComponentCache() + component = component_class(**default_kwargs) + component._shared_component_cache = cache + + with patch( + "lfx.components.data_source.sql_executor.SQLDatabase.from_uri", + return_value=created_db, + ) as mock_from_uri: + component.maybe_create_db() + + mock_from_uri.assert_called_once_with(default_kwargs["database_url"], engine_args=SQL_DATABASE_ENGINE_ARGS) + assert component.db is created_db + assert cache.values[default_kwargs["database_url"]] is created_db diff --git a/src/lfx/src/lfx/_assets/component_index.json b/src/lfx/src/lfx/_assets/component_index.json index 5059c70b1b..26aed68f3f 100644 --- a/src/lfx/src/lfx/_assets/component_index.json +++ b/src/lfx/src/lfx/_assets/component_index.json @@ -58688,7 +58688,7 @@ "icon": "database", "legacy": false, "metadata": { - "code_hash": "a8dd79af50b8", + "code_hash": "fb4189cb5af3", "dependencies": { "dependencies": [ { @@ -58772,7 +58772,7 @@ "show": true, "title_case": false, "type": "code", - "value": "from typing import TYPE_CHECKING, Any\n\nfrom langchain_community.utilities import SQLDatabase\nfrom sqlalchemy.exc import SQLAlchemyError\n\nfrom lfx.custom.custom_component.component_with_cache import ComponentWithCache\nfrom lfx.io import BoolInput, MessageTextInput, MultilineInput, Output\nfrom lfx.schema.dataframe import DataFrame\nfrom lfx.schema.message import Message\nfrom lfx.services.cache.utils import CacheMiss\n\nif TYPE_CHECKING:\n from sqlalchemy.engine import Result\n\n\nclass SQLComponent(ComponentWithCache):\n \"\"\"A sql component.\"\"\"\n\n display_name = \"SQL Database\"\n description = \"Executes SQL queries on SQLAlchemy-compatible databases.\"\n documentation: str = \"https://docs.langflow.org/sql-database\"\n icon = \"database\"\n name = \"SQLComponent\"\n metadata = {\"keywords\": [\"sql\", \"database\", \"query\", \"db\", \"fetch\"]}\n\n def __init__(self, **kwargs) -> None:\n super().__init__(**kwargs)\n self.db: SQLDatabase = None\n\n def maybe_create_db(self):\n if self.database_url != \"\":\n if self._shared_component_cache:\n cached_db = self._shared_component_cache.get(self.database_url)\n if not isinstance(cached_db, CacheMiss):\n self.db = cached_db\n return\n self.log(\"Connecting to database\")\n try:\n self.db = SQLDatabase.from_uri(self.database_url)\n except Exception as e:\n msg = f\"An error occurred while connecting to the database: {e}\"\n raise ValueError(msg) from e\n if self._shared_component_cache:\n self._shared_component_cache.set(self.database_url, self.db)\n\n inputs = [\n MessageTextInput(name=\"database_url\", display_name=\"Database URL\", required=True),\n MultilineInput(name=\"query\", display_name=\"SQL Query\", tool_mode=True, required=True),\n BoolInput(name=\"include_columns\", display_name=\"Include Columns\", value=True, tool_mode=True, advanced=True),\n BoolInput(\n name=\"add_error\",\n display_name=\"Add Error\",\n value=False,\n tool_mode=True,\n info=\"If True, the error will be added to the result\",\n advanced=True,\n ),\n ]\n\n outputs = [\n Output(display_name=\"Result Table\", name=\"run_sql_query\", method=\"run_sql_query\"),\n ]\n\n def build_component(\n self,\n ) -> Message:\n error = None\n self.maybe_create_db()\n try:\n result = self.db.run(self.query, include_columns=self.include_columns)\n self.status = result\n except SQLAlchemyError as e:\n msg = f\"An error occurred while running the SQL Query: {e}\"\n self.log(msg)\n result = str(e)\n self.status = result\n error = repr(e)\n\n if self.add_error and error is not None:\n result = f\"{result}\\n\\nError: {error}\\n\\nQuery: {self.query}\"\n elif error is not None:\n # Then we won't add the error to the result\n result = self.query\n\n return Message(text=result)\n\n def __execute_query(self) -> list[dict[str, Any]]:\n self.maybe_create_db()\n try:\n cursor: Result[Any] = self.db.run(self.query, fetch=\"cursor\")\n return [x._asdict() for x in cursor.fetchall()]\n except SQLAlchemyError as e:\n msg = f\"An error occurred while running the SQL Query: {e}\"\n self.log(msg)\n raise ValueError(msg) from e\n\n def run_sql_query(self) -> DataFrame:\n result = self.__execute_query()\n df_result = DataFrame(result)\n self.status = df_result\n return df_result\n" + "value": "from typing import TYPE_CHECKING, Any\n\nfrom langchain_community.utilities import SQLDatabase\nfrom sqlalchemy.exc import SQLAlchemyError\n\nfrom lfx.custom.custom_component.component_with_cache import ComponentWithCache\nfrom lfx.io import BoolInput, MessageTextInput, MultilineInput, Output\nfrom lfx.schema.dataframe import DataFrame\nfrom lfx.schema.message import Message\nfrom lfx.services.cache.utils import CacheMiss\n\nif TYPE_CHECKING:\n from sqlalchemy.engine import Result\n\nSQL_DATABASE_ENGINE_ARGS = {\"pool_pre_ping\": True}\n\n\nclass SQLComponent(ComponentWithCache):\n \"\"\"A sql component.\"\"\"\n\n display_name = \"SQL Database\"\n description = \"Executes SQL queries on SQLAlchemy-compatible databases.\"\n documentation: str = \"https://docs.langflow.org/sql-database\"\n icon = \"database\"\n name = \"SQLComponent\"\n metadata = {\"keywords\": [\"sql\", \"database\", \"query\", \"db\", \"fetch\"]}\n\n def __init__(self, **kwargs) -> None:\n super().__init__(**kwargs)\n self.db: SQLDatabase = None\n\n def maybe_create_db(self):\n if self.database_url != \"\":\n if self._shared_component_cache:\n cached_db = self._shared_component_cache.get(self.database_url)\n if not isinstance(cached_db, CacheMiss):\n self.db = cached_db\n return\n self.log(\"Connecting to database\")\n try:\n self.db = SQLDatabase.from_uri(self.database_url, engine_args=SQL_DATABASE_ENGINE_ARGS)\n except Exception as e:\n msg = f\"An error occurred while connecting to the database: {e}\"\n raise ValueError(msg) from e\n if self._shared_component_cache:\n self._shared_component_cache.set(self.database_url, self.db)\n\n inputs = [\n MessageTextInput(name=\"database_url\", display_name=\"Database URL\", required=True),\n MultilineInput(name=\"query\", display_name=\"SQL Query\", tool_mode=True, required=True),\n BoolInput(name=\"include_columns\", display_name=\"Include Columns\", value=True, tool_mode=True, advanced=True),\n BoolInput(\n name=\"add_error\",\n display_name=\"Add Error\",\n value=False,\n tool_mode=True,\n info=\"If True, the error will be added to the result\",\n advanced=True,\n ),\n ]\n\n outputs = [\n Output(display_name=\"Result Table\", name=\"run_sql_query\", method=\"run_sql_query\"),\n ]\n\n def build_component(\n self,\n ) -> Message:\n error = None\n self.maybe_create_db()\n try:\n result = self.db.run(self.query, include_columns=self.include_columns)\n self.status = result\n except SQLAlchemyError as e:\n msg = f\"An error occurred while running the SQL Query: {e}\"\n self.log(msg)\n result = str(e)\n self.status = result\n error = repr(e)\n\n if self.add_error and error is not None:\n result = f\"{result}\\n\\nError: {error}\\n\\nQuery: {self.query}\"\n elif error is not None:\n # Then we won't add the error to the result\n result = self.query\n\n return Message(text=result)\n\n def __execute_query(self) -> list[dict[str, Any]]:\n self.maybe_create_db()\n try:\n cursor: Result[Any] = self.db.run(self.query, fetch=\"cursor\")\n return [x._asdict() for x in cursor.fetchall()]\n except SQLAlchemyError as e:\n msg = f\"An error occurred while running the SQL Query: {e}\"\n self.log(msg)\n raise ValueError(msg) from e\n\n def run_sql_query(self) -> DataFrame:\n result = self.__execute_query()\n df_result = DataFrame(result)\n self.status = df_result\n return df_result\n" }, "database_url": { "_input_type": "MessageTextInput", @@ -118473,6 +118473,6 @@ "num_components": 354, "num_modules": 95 }, - "sha256": "836a9bd2862e2d4d276e3af6b9f89fc4c21f74bea8a2618b447906a8a48e0249", + "sha256": "1eaacdcac0e7f42362e192ca7a98b8e454abb57d636bb24f543b2599d10694db", "version": "1.10.1" } diff --git a/src/lfx/src/lfx/components/data_source/sql_executor.py b/src/lfx/src/lfx/components/data_source/sql_executor.py index c82bfa3760..ece37253a5 100644 --- a/src/lfx/src/lfx/components/data_source/sql_executor.py +++ b/src/lfx/src/lfx/components/data_source/sql_executor.py @@ -12,6 +12,8 @@ from lfx.services.cache.utils import CacheMiss if TYPE_CHECKING: from sqlalchemy.engine import Result +SQL_DATABASE_ENGINE_ARGS = {"pool_pre_ping": True} + class SQLComponent(ComponentWithCache): """A sql component.""" @@ -36,7 +38,7 @@ class SQLComponent(ComponentWithCache): return self.log("Connecting to database") try: - self.db = SQLDatabase.from_uri(self.database_url) + self.db = SQLDatabase.from_uri(self.database_url, engine_args=SQL_DATABASE_ENGINE_ARGS) except Exception as e: msg = f"An error occurred while connecting to the database: {e}" raise ValueError(msg) from e