From 2af2b80ba62c50e325d6500eedf7a9c2a5ee356b Mon Sep 17 00:00:00 2001
From: Deepesh Garg <42651287+deepeshgarg007@users.noreply.github.com>
Date: Thu, 14 Apr 2022 18:59:39 +0530
Subject: [PATCH] feat: Expose add index in exec_globals (#16605)

* feat: Expose add index in exec_globals

* fix: Linting issues
---
 .../server_script/test_server_script.py       | 22 +++++++++++++++++++
 frappe/utils/safe_exec.py                     |  3 +++
 2 files changed, 25 insertions(+)

diff --git a/frappe/core/doctype/server_script/test_server_script.py b/frappe/core/doctype/server_script/test_server_script.py
index 2685367695..fd600b8205 100644
--- a/frappe/core/doctype/server_script/test_server_script.py
+++ b/frappe/core/doctype/server_script/test_server_script.py
@@ -72,6 +72,16 @@ frappe.method_that_doesnt_exist("do some magic")
 		disabled=1,
 		script="""
 frappe.db.commit()
+""",
+	),
+	dict(
+		name="test_add_index",
+		script_type="DocType Event",
+		doctype_event="Before Save",
+		reference_doctype="ToDo",
+		disabled=1,
+		script="""
+frappe.db.add_index("Todo", ["color", "date"])
 """,
 	),
 ]
@@ -153,6 +163,18 @@ class TestServerScript(unittest.TestCase):
 		server_script.disabled = 1
 		server_script.save()
 
+	def test_add_index_in_doctype_event(self):
+		server_script = frappe.get_doc("Server Script", "test_add_index")
+		server_script.disabled = 0
+		server_script.save()
+
+		self.assertRaises(
+			AttributeError, frappe.get_doc(dict(doctype="ToDo", description="test me")).insert
+		)
+
+		server_script.disabled = 1
+		server_script.save()
+
 	def test_restricted_qb(self):
 		todo = frappe.get_doc(doctype="ToDo", description="QbScriptTestNote")
 		todo.insert()
diff --git a/frappe/utils/safe_exec.py b/frappe/utils/safe_exec.py
index 0a775fafec..aa6fa8b67f 100644
--- a/frappe/utils/safe_exec.py
+++ b/frappe/utils/safe_exec.py
@@ -56,8 +56,10 @@ def safe_exec(script, _globals=None, _locals=None, restrict_commit_rollback=Fals
 		exec_globals.update(_globals)
 
 	if restrict_commit_rollback:
+		# prevent user from using these in docevents
 		exec_globals.frappe.db.pop("commit", None)
 		exec_globals.frappe.db.pop("rollback", None)
+		exec_globals.frappe.db.pop("add_index", None)
 
 	# execute script compiled by RestrictedPython
 	frappe.flags.in_safe_exec = True
@@ -155,6 +157,7 @@ def get_safe_globals():
 				sql=read_sql,
 				commit=frappe.db.commit,
 				rollback=frappe.db.rollback,
+				add_index=frappe.db.add_index,
 			),
 		),
 		FrappeClient=FrappeClient,