Skip to content

Commit

Permalink
fix: Corrected URI handling in SQLDatabaseComponent (#3291)
Browse files Browse the repository at this point in the history
  • Loading branch information
carlosrcoelho authored Aug 12, 2024
1 parent 7264028 commit dc419d1
Showing 1 changed file with 8 additions and 5 deletions.
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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)

0 comments on commit dc419d1

Please sign in to comment.