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)