From dc419d11fdeebb815cf5f583a50b468b93a4af96 Mon Sep 17 00:00:00 2001 From: Carlos Coelho <80289056+carlosrcoelho@users.noreply.github.com> Date: Mon, 12 Aug 2024 19:45:13 -0300 Subject: [PATCH] fix: Corrected URI handling in SQLDatabaseComponent (#3291) --- .../components/langchain_utilities/SQLDatabase.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/backend/base/langflow/components/langchain_utilities/SQLDatabase.py b/src/backend/base/langflow/components/langchain_utilities/SQLDatabase.py index e9c7565d6e47..4d5104aabd00 100644 --- a/src/backend/base/langflow/components/langchain_utilities/SQLDatabase.py +++ b/src/backend/base/langflow/components/langchain_utilities/SQLDatabase.py @@ -1,6 +1,7 @@ -from langchain_experimental.sql.base import SQLDatabase - +from langchain_community.utilities.sql_database import SQLDatabase from langflow.custom import CustomComponent +from sqlalchemy import create_engine +from sqlalchemy.pool import StaticPool class SQLDatabaseComponent(CustomComponent): @@ -14,10 +15,12 @@ def build_config(self): } def clean_up_uri(self, uri: str) -> str: - if uri.startswith("postgresql://"): - uri = uri.replace("postgresql://", "postgres://") + if uri.startswith("postgres://"): + uri = uri.replace("postgres://", "postgresql://") return uri.strip() def build(self, uri: str) -> SQLDatabase: uri = self.clean_up_uri(uri) - return SQLDatabase.from_uri(uri) + # Create an engine using SQLAlchemy with StaticPool + engine = create_engine(uri, poolclass=StaticPool) + return SQLDatabase(engine)