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."""