fix: do not allow source.db access in script query

frappe · Nov 17, 2023 · 830c920 · 830c920
1 parent d56bbf8
commit 830c920
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/insights/insights/doctype/insights_data_source/insights_data_source.py b/insights/insights/doctype/insights_data_source/insights_data_source.py
@@ -64,6 +64,8 @@ def on_trash(self):
 
     @cached_property
     def db(self) -> BaseDatabase:
+        if frappe.flags.in_safe_exec:
+            raise NotImplementedError("Cannot access database in safe exec")
         if self.is_site_db:
             return SiteDB(data_source=self.name)
         if self.name == "Query Store":

diff --git a/insights/insights/doctype/insights_query/insights_script_query.py b/insights/insights/doctype/insights_query/insights_script_query.py
@@ -5,7 +5,7 @@
 import frappe
 import pandas as pd
 from frappe.utils.password import get_decrypted_password
-from frappe.utils.safe_exec import compile_restricted, get_safe_globals
+from frappe.utils.safe_exec import safe_exec
 
 from insights import notify
 
@@ -45,10 +45,11 @@ def get_value(variable):
             variables = self.doc.get("variables") or []
             variables = {var.variable_name: get_value(var) for var in variables}
             _locals = {"results": results, **variables}
-            exec(
-                compile_restricted(script, filename="<scriptquery>"),
-                get_safe_exec_globals(),
-                _locals,
+            safe_exec(
+                script,
+                _globals=get_globals(),
+                _locals=_locals,
+                restrict_commit_rollback=True,
             )
             self.update_script_log()
             results = _locals["results"]
@@ -107,16 +108,23 @@ def get_selected_tables(self):
         return []
 
 
-def get_safe_exec_globals():
-    safe_globals = get_safe_globals()
-
+def get_globals():
     pandas = frappe._dict()
     pandas.DataFrame = pd.DataFrame
     pandas.read_csv = pd.read_csv
     pandas.json_normalize = pd.json_normalize
     # mock out to_csv and to_json to prevent users from writing to disk
     pandas.DataFrame.to_csv = lambda *args, **kwargs: None
     pandas.DataFrame.to_json = lambda *args, **kwargs: None
-    safe_globals.pandas = pandas
 
-    return safe_globals
+    return {
+        "pandas": pandas,
+        "get_query_results": get_query_results,
+    }
+
+
+def get_query_results(query_name):
+    if not isinstance(query_name, str):
+        raise ScriptQueryExecutionError("Query name should be a string.")
+    doc = frappe.get_doc("Insights Query", query_name)
+    return doc.retrieve_results(fetch_if_not_cached=True)