From b6ff3e5ddd5ee6c3e4b24365eb653c313f821e9b Mon Sep 17 00:00:00 2001
From: wlwlwlzhang <wlwlwlzhang@tencent.com>
Date: Tue, 5 Sep 2023 18:13:15 +0800
Subject: [PATCH] Close the connection  when using a context manager

---
 dataset/database.py | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/dataset/database.py b/dataset/database.py
index d8a07ad..467a83e 100644
--- a/dataset/database.py
+++ b/dataset/database.py
@@ -158,15 +158,22 @@ def __enter__(self):
         return self
 
     def __exit__(self, error_type, error_value, traceback):
-        """End a transaction by committing or rolling back."""
-        if error_type is None:
-            try:
-                self.commit()
-            except Exception:
-                with safe_reraise():
-                    self.rollback()
-        else:
-            self.rollback()
+        """End a transaction by committing or rolling back. Close local connection"""
+        try:
+            if error_type is None:
+                try:
+                    self.commit()
+                except Exception:
+                    with safe_reraise():
+                        self.rollback()
+            else:
+                self.rollback()
+        except Exception:
+            raise
+        finally:
+            tx_conn = self.connections.pop(threading.get_ident(), None)
+            if tx_conn is not None:
+                tx_conn.close()
 
     def close(self):
         """Close database connections. Makes this object unusable."""