fix: handle interface error during report timeout (#25893)

frappe/core/doctype/prepared_report/prepared_report.py

@@ -8,6 +8,7 @@ from typing import Any
 from rq import get_current_job
 
 import frappe
+from frappe.database.utils import dangerously_reconnect_on_connection_abort
 from frappe.desk.form.load import get_attachments
 from frappe.desk.query_report import generate_report_result
 from frappe.model.document import Document
@@ -115,6 +116,7 @@ def generate_report(prepared_report):
 	except Exception:
 		instance.status = "Error"
 		instance.error_message = frappe.get_traceback(with_context=True)
+		_save_instance(instance)  # we need to ensure that error gets stored
 
 	instance.report_end_time = frappe.utils.now()
 	instance.save(ignore_permissions=True)
@@ -126,6 +128,11 @@ def generate_report(prepared_report):
 	)
 
 
+@dangerously_reconnect_on_connection_abort
+def _save_instance(instance):
+	instance.save(ignore_permissions=True)
+
+
 def update_job_id(prepared_report):
 	job = get_current_job()
 

frappe/database/database.py

@@ -1245,7 +1245,7 @@ class Database:
 
 	@staticmethod
 	def is_column_missing(e):
-		return frappe.db.is_missing_column(e)
+		raise NotImplementedError
 
 	def get_descendants(self, doctype, name):
 		"""Return descendants of the group node in tree"""

frappe/database/mariadb/database.py

@@ -97,6 +97,10 @@ class MariaDBExceptionUtil:
 			and isinstance(e, pymysql.IntegrityError)
 		)
 
+	@staticmethod
+	def is_interface_error(e: pymysql.Error):
+		return isinstance(e, pymysql.InterfaceError)
+
 
 class MariaDBConnectionUtil:
 	def get_connection(self):

frappe/database/postgres/database.py

@@ -13,6 +13,7 @@ from psycopg2.errorcodes import (
 	UNIQUE_VIOLATION,
 )
 from psycopg2.errors import (
+	InterfaceError,
 	LockNotAvailable,
 	ReadOnlySqlTransaction,
 	SequenceGeneratorLimitExceeded,
@@ -116,6 +117,10 @@ class PostgresExceptionUtil:
 	def is_db_table_size_limit(e) -> bool:
 		return False
 
+	@staticmethod
+	def is_interface_error(e):
+		return isinstance(e, InterfaceError)
+
 
 class PostgresDatabase(PostgresExceptionUtil, Database):
 	REGEX_CHARACTER = "~"

frappe/database/utils.py

@@ -1,15 +1,14 @@
 # Copyright (c) 2022, Frappe Technologies Pvt. Ltd. and Contributors
 # License: MIT. See LICENSE
 
-from functools import cached_property
-from types import NoneType
+from functools import cached_property, wraps
 
 import frappe
 from frappe.query_builder.builder import MariaDB, Postgres
 from frappe.query_builder.functions import Function
 
 Query = str | MariaDB | Postgres
-QueryValues = tuple | list | dict | NoneType
+QueryValues = tuple | list | dict | None
 
 EmptyQueryValues = object()
 FallBackDateTimeStr = "0001-01-01 00:00:00.000000"
@@ -72,3 +71,26 @@ class LazyMogrify(LazyString):
 
 	def _setup(self) -> str:
 		return frappe.db.mogrify(self.query, self.values)
+
+
+def dangerously_reconnect_on_connection_abort(func):
+	"""Reconnect on connection failure.
+
+	As the name suggest, it's dangerous to use this function as it will NOT restore DB transaction
+	so make sure you're using it right.
+
+	Ideal use case: Some kinda logging or final steps in a background jobs. Anything more than that
+	will risk bugs from DB transactions.
+	"""
+
+	@wraps(func)
+	def wrapper(*args, **kwargs):
+		try:
+			return func(*args, **kwargs)
+		except Exception as e:
+			if frappe.db.is_interface_error(e):
+				frappe.db.connect()
+				return func(*args, **kwargs)
+			raise
+
+	return wrapper