pengjunjie/langflow
1
0
Fork 1
mirror of https://github.com/langflow-ai/langflow.git synced 2026-07-24 13:09:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: Corrected URI handling in SQLDatabaseComponent (#3291)

Browse Source
This commit is contained in:
Carlos Coelho
2024-08-12 19:45:13 -03:00
committed by GitHub
parent 7264028e41
commit dc419d11fd
1 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/backend/base/langflow/components/langchain_utilities/SQLDatabase.py
View File

@ -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 @@ class SQLDatabaseComponent(CustomComponent):
}
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)