refactor: code duplication for sql whitelisting

frappe/core/doctype/report/report.py

@@ -11,7 +11,7 @@ from frappe.modules import make_boilerplate
 from frappe.core.doctype.page.page import delete_custom_role
 from frappe.core.doctype.custom_role.custom_role import get_custom_allowed_roles
 from frappe.desk.reportview import append_totals_row
-from frappe.utils.safe_exec import safe_exec
+from frappe.utils.safe_exec import safe_exec, check_safe_sql_query
 
 
 class Report(Document):
@@ -110,15 +110,7 @@ class Report(Document):
 		if not self.query:
 			frappe.throw(_("Must specify a Query to run"), title=_('Report Document Error'))
 
-		# Disallow SQL that writes to the database.
-		if (not self.query.lower().startswith("select") and
-		    not self.query.lower().startswith("with")):
-			frappe.throw(_("Query must be a SELECT or WITH"), title=_('Report Document Error'))
-
-		# As of MariaDB 10.9, CTE WITH statements can only be combined with a SELECT clause and
-		# therefore are read-only. Postgres allows WITH ... INSERT INTO statements.
-		if (self.query.lower().startswith("with") and frappe.db.db_type != "mariadb"):
-			frappe.throw(_("WITH queries are only allowed for MariaDB databases"), title=_('Report Document Error'))
+		check_safe_sql_query(self.query)
 
 		result = [list(t) for t in frappe.db.sql(self.query, filters)]
 		columns = self.get_columns() or [cstr(c[0]) for c in frappe.db.get_description()]

frappe/core/doctype/report/test_report.py

@@ -1,17 +1,19 @@
 # Copyright (c) 2015, Frappe Technologies Pvt. Ltd. and Contributors
 # License: MIT. See LICENSE
 
+import textwrap
+
 import frappe, json, os
-import unittest
 from frappe.desk.query_report import run, save_report, add_total_row
 from frappe.desk.reportview import delete_report, save_report as _save_report
 from frappe.custom.doctype.customize_form.customize_form import reset_customization
 from frappe.core.doctype.user_permission.test_user_permission import create_user
+from frappe.tests.utils import FrappeTestCase
 
 test_records = frappe.get_test_records('Report')
 test_dependencies = ['User']
 
-class TestReport(unittest.TestCase):
+class TestReport(FrappeTestCase):
 	def test_report_builder(self):
 		if frappe.db.exists('Report', 'User Activity Report'):
 			frappe.delete_doc('Report', 'User Activity Report')
@@ -335,3 +337,29 @@ result = [
 		self.assertEqual(result[-1][0], "Total")
 		self.assertEqual(result[-1][1], 200)
 		self.assertEqual(result[-1][2], 150.50)
+
+	def test_cte_in_query_report(self):
+		cte_query = textwrap.dedent("""
+			with enabled_users as (
+				select name
+				from `tabUser`
+				where enabled = 1
+			)
+			select * from enabled_users;
+		""")
+
+		report = frappe.get_doc({
+			"doctype": "Report",
+			"ref_doctype": "User",
+			"report_name": "Enabled Users List",
+			"report_type": "Query Report",
+			"is_standard": "No",
+			"query": cte_query,
+		}).insert()
+
+		if frappe.db.db_type == "mariadb":
+			col, rows = report.execute_query_report(filters={})
+			self.assertEqual(col[0], "name")
+			self.assertGreaterEqual(len(rows), 1)
+		elif frappe.db.db_type == "postgres":
+			self.assertRaises(frappe.PermissionError, report.execute_query_report, filters={})

frappe/query_builder/utils.py

@@ -54,6 +54,9 @@ def patch_query_execute():
 	This excludes the use of `frappe.db.sql` method while
 	executing the query object
 	"""
+	from frappe.utils.safe_exec import check_safe_sql_query
+
+
 	def execute_query(query, *args, **kwargs):
 		query, params = prepare_query(query)
 		return frappe.db.sql(query, params, *args, **kwargs) # nosemgrep
@@ -63,7 +66,7 @@ def patch_query_execute():
 
 		param_collector = NamedParameterWrapper()
 		query = query.get_sql(param_wrapper=param_collector)
-		if frappe.flags.in_safe_exec and not query.lower().strip().startswith("select"):
+		if frappe.flags.in_safe_exec and not check_safe_sql_query(query, throw=False):
 			callstack = inspect.stack()
 			if len(callstack) >= 3 and ".py" in callstack[2].filename:
 				# ignore any query builder methods called from python files
@@ -77,7 +80,7 @@ def patch_query_execute():
 				#
 				# if frame2 is server script it wont have a filename and hence
 				# it shouldn't be allowed.
-				# ps. stack() returns `"<unknown>"` as filename.
+				# p.s. stack() returns `"<unknown>"` as filename if not a file.
 				pass
 			else:
 				raise frappe.PermissionError('Only SELECT SQL allowed in scripting')

frappe/utils/safe_exec.py

@@ -269,11 +269,33 @@ def get_hooks(hook=None, default=None, app_name=None):
 def read_sql(query, *args, **kwargs):
 	'''a wrapper for frappe.db.sql to allow reads'''
 	query = str(query)
-	if frappe.flags.in_safe_exec and not query.strip().lower().startswith('select'):
-		raise frappe.PermissionError('Only SELECT SQL allowed in scripting')
+	if frappe.flags.in_safe_exec:
+		check_safe_sql_query(query)
 	return frappe.db.sql(query, *args, **kwargs)
 
 
+def check_safe_sql_query(query: str, throw: bool = True) -> bool:
+	""" Check if SQL query is safe for running in restricted context.
+
+		Safe queries:
+			1. Read only 'select' or 'explain' queries
+			2. CTE on mariadb where writes are not allowed.
+	"""
+
+	query = query.strip().lower()
+	whitelisted_statements = ("select", "explain")
+
+	if (query.startswith(whitelisted_statements)
+		or (query.startswith("with") and frappe.db.db_type == "mariadb")):
+		return True
+
+	if throw:
+		frappe.throw(_("Query must be of SELECT or read-only WITH type."),
+				title=_("Unsafe SQL query"), exc=frappe.PermissionError)
+
+	return False
+
+
 def _getitem(obj, key):
 	# guard function for RestrictedPython
 	# allow any key to be accessed as long as it does not start with underscore