restricted sql permissions to SELECT/PRAGMA queries only

Author: codinglabsong

src/any_chatbot/tools.py

@@ -40,11 +40,29 @@ def retrieve(
     return retrieve
 
 
+def is_safe_sql(query: str) -> bool:
+    """Filters out destructive sql queries"""
+    forbidden = ["insert", "update", "delete", "drop", "alter", "create", "replace"]
+    # Make sure to only block whole words (e.i., don't block 'updated_at')
+    return not any(f" {word} " in f" {query.lower()} " for word in forbidden)
+
+
 def initialize_sql_toolkit(
     llm,
-    db_path: Path = DATA / "csv_excel_to_db" / "my_data.duckdb",
+    db_path: Path = DATA / "generated_db" / "csv_excel_to_db.duckdb",
 ):
     db = SQLDatabase.from_uri(f"duckdb:///{db_path}")
+
+    # Monkey-path the run method to include safety filter
+    original_run = db.run
+
+    def safe_run(query: str, *args, **kwargs):
+        if not is_safe_sql(query):
+            return "Query blocked: Only SELECT/PRAGMA queries are allowed."
+        return original_run(query, *args, **kwargs)
+
+    db.run = safe_run
+
     toolkit = SQLDatabaseToolkit(db=db, llm=llm)
     tools = toolkit.get_tools()
     return tools