fix(component): reconnect stale cached SQL database connections on SQL Database component (#13733)

* Implement check for stale cached database connection

Add stale connection check for cached database.

* [autofix.ci] apply automated fixes

* fix(component): use SQLAlchemy pre-ping for SQL database cache

* [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>

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

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"
 }

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