From fefd9ac2e2190d37d3669390a2d6285506a2646c Mon Sep 17 00:00:00 2001
From: Ankush Menat <ankush@frappe.io>
Date: Fri, 14 Jul 2023 13:56:57 +0530
Subject: [PATCH] fix: `connect_replica` and `read_only` should be idempotent
 (#21674)

* fix: `connect_replica` should be idempotent

Calling `connect_replica` twice ends up forgetting orginal writable
replica completely.

* test: replicas

closes https://github.com/frappe/frappe/issues/21619
---
 frappe/__init__.py      | 15 ++++++++++++---
 frappe/tests/test_db.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/frappe/__init__.py b/frappe/__init__.py
index 6f1223a748..c3c5a107d7 100644
--- a/frappe/__init__.py
+++ b/frappe/__init__.py
@@ -274,9 +274,12 @@ def connect(
 		set_user("Administrator")
 
 
-def connect_replica():
+def connect_replica() -> bool:
 	from frappe.database import get_db
 
+	if local and hasattr(local, "replica_db") and hasattr(local, "primary_db"):
+		return False
+
 	user = local.conf.db_name
 	password = local.conf.db_password
 	port = local.conf.replica_db_port
@@ -291,6 +294,8 @@ def connect_replica():
 	local.primary_db = local.db
 	local.db = local.replica_db
 
+	return True
+
 
 def get_site_config(sites_path: str | None = None, site_path: str | None = None) -> dict[str, Any]:
 	"""Returns `site_config.json` combined with `sites/common_site_config.json`.
@@ -809,13 +814,17 @@ def is_whitelisted(method):
 def read_only():
 	def innfn(fn):
 		def wrapper_fn(*args, **kwargs):
+
+			# frappe.read_only could be called from nested functions, in such cases don't swap the
+			# connection again.
+			switched_connection = False
 			if conf.read_from_replica:
-				connect_replica()
+				switched_connection = connect_replica()
 
 			try:
 				retval = fn(*args, **get_newargs(fn, kwargs))
 			finally:
-				if local and hasattr(local, "primary_db"):
+				if switched_connection and local and hasattr(local, "primary_db"):
 					local.db.close()
 					local.db = local.primary_db
 
diff --git a/frappe/tests/test_db.py b/frappe/tests/test_db.py
index e8afba5d96..76722fccc7 100644
--- a/frappe/tests/test_db.py
+++ b/frappe/tests/test_db.py
@@ -930,3 +930,34 @@ class TestTransactionManagement(FrappeTestCase):
 
 		frappe.db.commit()
 		self.assertEqual(_get_transaction_id(), _get_transaction_id())
+
+
+# Treat same DB as replica for tests, a separate connection will be opened
+class TestReplicaConnections(FrappeTestCase):
+	def test_switching_to_replica(self):
+		with patch.dict(frappe.local.conf, {"read_from_replica": 1, "replica_host": "localhost"}):
+
+			def db_id():
+				return id(frappe.local.db)
+
+			write_connection = db_id()
+			read_only_connection = None
+
+			@frappe.read_only()
+			def outer():
+				nonlocal read_only_connection
+				read_only_connection = db_id()
+
+				# A new connection should be opened
+				self.assertNotEqual(read_only_connection, write_connection)
+				inner()
+				# calling nested read only function shouldn't change connection
+				self.assertEqual(read_only_connection, db_id())
+
+			@frappe.read_only()
+			def inner():
+				# calling nested read only function shouldn't change connection
+				self.assertEqual(read_only_connection, db_id())
+
+			outer()
+			self.assertEqual(write_connection, db_id())