Skip to content

Commit

Permalink
readd my stuff
Browse files Browse the repository at this point in the history
  • Loading branch information
Uranium2 committed Feb 18, 2024
1 parent 9077ec6 commit 6feba54
Showing 1 changed file with 17 additions and 5 deletions.
22 changes: 17 additions & 5 deletions libs/community/langchain_community/utilities/sql_database.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ def __init__(
custom_table_info: Optional[dict] = None,
view_support: bool = False,
max_string_length: int = 300,
restricted_keywords: Optional[List[str]] = None,
):
"""Create engine from database URI."""
self._engine = engine
Expand Down Expand Up @@ -123,6 +124,10 @@ def __init__(
schema=self._schema,
)

# Restricted keywords to not execute on database
self.restricted_keywords = restricted_keywords if restricted_keywords else []


@classmethod
def from_uri(
cls, database_uri: str, engine_args: Optional[dict] = None, **kwargs: Any
Expand Down Expand Up @@ -447,11 +452,18 @@ def _execute(
pass
else:
raise TypeError(f"Query expression has unknown type: {type(command)}")
cursor = connection.execute(
command,
parameters,
execution_options=execution_options,
)
if not self.detect_restricted_keywords(command):
cursor = connection.execute(
command,
parameters,
execution_options=execution_options,
)
else:
raise PermissionError(
f"Restricted keywords in the SQL '{command}' "
f"Commands '{self.restricted_keywords}' are forbidden."
)


if cursor.returns_rows:
if fetch == "all":
Expand Down

0 comments on commit 6feba54

Please sign in to comment.