refactor(safe_exec): Manage in-safe frappe.db.sql

* Check if flag in_safe_exec to check if SELECT query in frappe.db.sql
* Get rid of read_sql...refactor logic
* Enable frappe.qb just like that ;)
* Add support fro *args in .run

frappe/database/database.py

@@ -109,7 +109,13 @@ class Database(object):
 				{"name": "a%", "owner":"test@example.com"})
 
 		"""
+
 		query = str(query)
+
+		if frappe.flags.in_safe_exec:
+			if not query.strip().lower().startswith('select'):
+				raise frappe.PermissionError('Only SELECT SQL allowed in scripting')
+
 		if not run:
 			return query
 		if re.search(r'ifnull\(', query, flags=re.IGNORECASE):

frappe/query_builder/builder.py

@@ -2,6 +2,7 @@ from pypika import MySQLQuery, Order, PostgreSQLQuery, terms
 from pypika.queries import Schema, Table
 from frappe.utils import get_table_name
 from pypika.terms import Function
+
 class Base:
 	terms = terms
 	desc = Order.desc

frappe/query_builder/utils.py

@@ -50,8 +50,8 @@ def patch_query_execute():
 	executing the query object
 	"""
 
-	def execute_query(query, **kwargs):
-		return frappe.db.sql(query, **kwargs)
+	def execute_query(query, *args, **kwargs):
+		return frappe.db.sql(query, *args, **kwargs)
 
 	query_class = get_attr(str(frappe.qb).split("'")[1])
 	builder_class = get_type_hints(query_class._builder).get('return')

frappe/utils/safe_exec.py

@@ -1,25 +1,28 @@
 
-import os, json, inspect
+import inspect
+import json
 import mimetypes
+
+import RestrictedPython.Guards
 from html2text import html2text
 from RestrictedPython import compile_restricted, safe_globals
-import RestrictedPython.Guards
+
 import frappe
-from frappe import _
-import frappe.utils
-import frappe.utils.data
-from frappe.website.utils import (get_shade, get_toc, get_next_link)
-from frappe.modules import scrub
-from frappe.www.printview import get_visible_columns
 import frappe.exceptions
 import frappe.integrations.utils
+import frappe.utils
+import frappe.utils.data
+from frappe import _
 from frappe.frappeclient import FrappeClient
-from frappe.query_builder.utils import get_attr
-from typing import get_type_hints
+from frappe.modules import scrub
+from frappe.website.utils import get_next_link, get_shade, get_toc
+from frappe.www.printview import get_visible_columns
+
 
 class ServerScriptNotEnabled(frappe.PermissionError):
 	pass
 
+
 class NamespaceDict(frappe._dict):
 	"""Raise AttributeError if function not found in namespace"""
 	def __getattr__(self, key):
@@ -30,18 +33,6 @@ class NamespaceDict(frappe._dict):
 			return default_function
 		return ret
 
-def get_safe_query_builder():
-	"""Allows execution of SELECT SQL queries only.
-	Raises:
-	   	PermissionsError raised on execution of any other SQL query
-	"""
-	query_class = get_attr(str(frappe.qb).split("'")[1])
-	class SafeQB(query_class):
-		def __init__(self, *args, **kwargs):
-			_builder = get_type_hints(super()._builder).get('return')
-			_builder.run = read_sql
-
-	return SafeQB()
 
 def safe_exec(script, _globals=None, _locals=None):
 	# server scripts can be disabled via site_config.json
@@ -60,13 +51,15 @@ def safe_exec(script, _globals=None, _locals=None):
 		exec_globals.update(_globals)
 
 	# execute script compiled by RestrictedPython
+	frappe.flags.in_safe_exec = True
 	exec(compile_restricted(script), exec_globals, _locals) # pylint: disable=exec-used
+	frappe.flags.in_safe_exec = False
 
 	return exec_globals, _locals
 
 def get_safe_globals():
 	datautils = frappe._dict()
-	safe_qb = get_safe_query_builder()
+
 	if frappe.db:
 		date_format = frappe.db.get_default("date_format") or "yyyy-mm-dd"
 		time_format = frappe.db.get_default("time_format") or "HH:mm:ss"
@@ -100,7 +93,7 @@ def get_safe_globals():
 			bold=frappe.bold,
 			copy_doc=frappe.copy_doc,
 			errprint=frappe.errprint,
-			qb=safe_qb,
+			qb=frappe.qb,
 
 			get_meta=frappe.get_meta,
 			get_doc=frappe.get_doc,
@@ -169,7 +162,7 @@ def get_safe_globals():
 			avg=frappe.db.avg,
 			sum=frappe.db.sum,
 			escape=frappe.db.escape,
-			sql=read_sql
+			sql=frappe.db.sql
 		)
 
 	if frappe.response:
@@ -189,14 +182,6 @@ def get_safe_globals():
 
 	return out
 
-def read_sql(query, *args, **kwargs):
-	'''a wrapper for frappe.db.sql to allow reads'''
-	query = str(query)
-	if query.strip().split(None, 1)[0].lower() == 'select':
-		return frappe.db.sql(query, *args, **kwargs)
-	else:
-		raise frappe.PermissionError('Only SELECT SQL allowed in scripting')
-
 def run_script(script):
 	'''run another server script'''
 	return frappe.get_doc('Server Script', script).execute_method()