From ec3f705e4f3087ab0c77ba5bc1fbe03f8c1c0899 Mon Sep 17 00:00:00 2001
From: Ankush Menat <ankush@frappe.io>
Date: Fri, 9 Dec 2022 14:18:49 +0530
Subject: [PATCH] feat: finer frappe Recorder control with decorator (#19220)

Currently frappe recorder can be enabled globally and profiles every
request. This is often way too much info. If you already know where
problem lies you use this decorator sparingly to only profile relevant
functions.

Usage:
```py
from frappe.recorder import record_queries

@record_queries
def sus_slow_function():
    frappe.db.sql("select everything from everywhere")
```
---
 frappe/recorder.py            | 44 +++++++++++++++++++++++++++++------
 frappe/tests/test_recorder.py | 12 ++++++++++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/frappe/recorder.py b/frappe/recorder.py
index b00f15c6b5..537d1ee996 100644
--- a/frappe/recorder.py
+++ b/frappe/recorder.py
@@ -1,11 +1,13 @@
 # Copyright (c) 2018, Frappe Technologies Pvt. Ltd. and Contributors
 # License: MIT. See LICENSE
 import datetime
+import functools
 import inspect
 import json
 import re
 import time
 from collections import Counter
+from typing import Callable
 
 import sqlparse
 
@@ -61,9 +63,9 @@ def get_current_stack_frames():
 		pass
 
 
-def record():
+def record(force=False):
 	if __debug__:
-		if frappe.cache().get_value(RECORDER_INTERCEPT_FLAG):
+		if frappe.cache().get_value(RECORDER_INTERCEPT_FLAG) or force:
 			frappe.local._recorder = Recorder()
 
 
@@ -78,11 +80,19 @@ class Recorder:
 		self.uuid = frappe.generate_hash(length=10)
 		self.time = datetime.datetime.now()
 		self.calls = []
-		self.path = frappe.request.path
-		self.cmd = frappe.local.form_dict.cmd or ""
-		self.method = frappe.request.method
-		self.headers = dict(frappe.local.request.headers)
-		self.form_dict = frappe.local.form_dict
+		if frappe.request:
+			self.path = frappe.request.path
+			self.cmd = frappe.local.form_dict.cmd or ""
+			self.method = frappe.request.method
+			self.headers = dict(frappe.local.request.headers)
+			self.form_dict = frappe.local.form_dict
+		else:
+			self.path = None
+			self.cmd = None
+			self.method = None
+			self.headers = None
+			self.form_dict = None
+
 		_patch()
 
 	def register(self, data):
@@ -125,6 +135,10 @@ def _patch():
 	frappe.db.sql = sql
 
 
+def _unpatch():
+	frappe.db.sql = frappe.db._sql
+
+
 def do_not_record(function):
 	def wrapper(*args, **kwargs):
 		if hasattr(frappe.local, "_recorder"):
@@ -189,3 +203,19 @@ def export_data(*args, **kwargs):
 def delete(*args, **kwargs):
 	frappe.cache().delete_value(RECORDER_REQUEST_SPARSE_HASH)
 	frappe.cache().delete_value(RECORDER_REQUEST_HASH)
+
+
+def record_queries(func: Callable):
+	"""Decorator to profile a specific function using recorder."""
+
+	@functools.wraps(func)
+	def wrapped(*args, **kwargs):
+		record(force=True)
+		frappe.local._recorder.path = f"Function call: {func.__module__}.{func.__qualname__}"
+		ret = func(*args, **kwargs)
+		dump()
+		_unpatch()
+		print("Recorded queries, open recorder to view them.")
+		return ret
+
+	return wrapped
diff --git a/frappe/tests/test_recorder.py b/frappe/tests/test_recorder.py
index a265ce9b1b..ffe3df4b23 100644
--- a/frappe/tests/test_recorder.py
+++ b/frappe/tests/test_recorder.py
@@ -123,3 +123,15 @@ class TestRecorder(FrappeTestCase):
 	def test_error_page_rendering(self):
 		content = get_response_content("error")
 		self.assertIn("Error", content)
+
+
+class TestRecorderDeco(FrappeTestCase):
+	def test_recorder_flag(self):
+		frappe.recorder.delete()
+
+		@frappe.recorder.record_queries
+		def test():
+			frappe.get_all("User")
+
+		test()
+		self.assertTrue(frappe.recorder.get())