From b35a55ed542127756ac75543ffc0c7d9d078bdc4 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Sat, 8 Feb 2025 10:36:01 +0800
Subject: [PATCH] fix: Cannot set original org when exception occurs

---
 apps/orgs/utils.py | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/apps/orgs/utils.py b/apps/orgs/utils.py
index 26c5a408a..21f96f20c 100644
--- a/apps/orgs/utils.py
+++ b/apps/orgs/utils.py
@@ -97,20 +97,24 @@ def get_current_org_id_for_serializer():
 @contextmanager
 def tmp_to_root_org():
     ori_org = get_current_org()
-    set_to_root_org()
-    yield
-    if ori_org is not None:
-        set_current_org(ori_org)
+    try:
+        set_to_root_org()
+        yield
+    finally:
+        if ori_org is not None:
+            set_current_org(ori_org)
 
 
 @contextmanager
 def tmp_to_org(org):
     ori_org = get_current_org()
-    if org:
-        set_current_org(org)
-    yield
-    if ori_org is not None:
-        set_current_org(ori_org)
+    try:
+        if org:
+            set_current_org(org)
+        yield
+    finally:
+        if ori_org is not None:
+            set_current_org(ori_org)
 
 
 @contextmanager
@@ -122,9 +126,10 @@ def tmp_to_builtin_org(system=0, default=0):
     else:
         raise ValueError("Must set system or default")
     ori_org = get_current_org()
-    set_current_org(org_id)
-    yield
-    if ori_org is not None:
+    try:
+        set_current_org(org_id)
+        yield
+    finally:
         set_current_org(ori_org)