Skip to content
This repository was archived by the owner on May 17, 2024. It is now read-only.

Commit af36b66

Browse files
author
Sung Won Chung
committed
resync
1 parent 73c7901 commit af36b66

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

data_diff/databases/base.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -922,6 +922,12 @@ class Database(abc.ABC):
922922
is_closed: bool = False
923923
_dialect: BaseDialect = None
924924

925+
def __enter__(self):
926+
return self
927+
928+
def __exit__(self, exc_type, exc_value, traceback):
929+
self.close()
930+
925931
@property
926932
def name(self):
927933
return type(self).__name__
@@ -1164,9 +1170,8 @@ def _query_conn(self, conn, sql_code: Union[str, ThreadLocalInterpreter]) -> Que
11641170
return apply_query(callback, sql_code)
11651171

11661172
def close(self):
1167-
"Close connection(s) to the database instance. Querying will stop functioning."
1173+
"""Close connection(s) to the database instance. Querying will stop functioning."""
11681174
self.is_closed = True
1169-
return super().close()
11701175

11711176
@property
11721177
def dialect(self) -> BaseDialect:
@@ -1237,6 +1242,8 @@ def create_connection(self):
12371242
def close(self):
12381243
super().close()
12391244
self._queue.shutdown()
1245+
if hasattr(self.thread_local, "conn"):
1246+
self.thread_local.conn.close()
12401247

12411248
@property
12421249
def is_autocommit(self) -> bool:
@@ -1263,4 +1270,4 @@ def is_autocommit(self) -> bool:
12631270
DEFAULT_DATETIME_PRECISION = 6
12641271
DEFAULT_NUMERIC_PRECISION = 24
12651272

1266-
TIMESTAMP_PRECISION_POS = 20 # len("2022-06-03 12:24:35.") == 20
1273+
TIMESTAMP_PRECISION_POS = 20 # len("2022-06-03 12:24:35.") == 20

0 commit comments

Comments
 (0)