Add a close method for Snowflake, Databricks, Presto and BigQuery classes

Author: pik94

data_diff/database.py

@@ -132,6 +132,11 @@ def _query(self, sql_code: str) -> list:
         "Send query to database and return result"
         ...
 
+    @abstractmethod
+    def close(self):
+        """Close a connection"""
+        ...
+
 
 class ThreadedDatabase(Database):
     """Access the database through singleton threads.
@@ -213,6 +218,9 @@ def _query(self, sql_code: str) -> list:
         "Uses the standard SQL cursor interface"
         return _query_conn(self._conn, sql_code)
 
+    def close(self):
+        self._conn.close()
+
 
 class MySQL(ThreadedDatabase):
     def __init__(self, host, port, database, user, password, *, thread_count):
@@ -333,6 +341,9 @@ def _query(self, sql_code: str):
     def to_string(self, s: str):
         return f"cast({s} as string)"
 
+    def close(self):
+        self._client.close()
+
 
 class Snowflake(Database):
     def __init__(
@@ -377,6 +388,9 @@ def md5_to_int(self, s: str) -> str:
     def to_string(self, s: str):
         return f"cast({s} as string)"
 
+    def close(self):
+        self._conn.close()
+
 
 class Databricks(Database):
     def __init__(
@@ -404,6 +418,9 @@ def md5_to_int(self, s: str) -> str:
     def to_string(self, s: str) -> str:
         return f"cast({s} as string)"
 
+    def close(self):
+        self._conn.close()
+
 
 def connect_to_uri(db_uri: str, thread_count: Optional[int] = 1) -> Database:
     """Connect to the given database uri