fix: Close SQL Alchemy connections. (#453)

This addresses #868, where extractions using the Snowflake extractor would
hand until the connection was closed. I've added close() calls to the other
SQL Alchemy based extractors since it is a general expectation of SQL Alchemy
and most of the infrastructure existed to do it already.

Signed-off-by: Dave Cameron <dcameron@digitalocean.com>

databuilder/extractor/athena_metadata_extractor.py

@@ -63,6 +63,9 @@ def init(self, conf: ConfigTree) -> None:
         self._alchemy_extractor.init(sql_alch_conf)
         self._extract_iter: Union[None, Iterator] = None
 
+    def close(self) -> None:
+        self._alchemy_extractor.close()
+
     def extract(self) -> Union[TableMetadata, None]:
         if not self._extract_iter:
             self._extract_iter = self._get_extract_iter()

databuilder/extractor/druid_metadata_extractor.py

@@ -59,6 +59,9 @@ def init(self, conf: ConfigTree) -> None:
         self._alchemy_extractor.init(sql_alch_conf)
         self._extract_iter: Union[None, Iterator] = None
 
+    def close(self) -> None:
+        self._alchemy_extractor.close()
+
     def extract(self) -> Union[TableMetadata, None]:
         if not self._extract_iter:
             self._extract_iter = self._get_extract_iter()

databuilder/extractor/mssql_metadata_extractor.py

@@ -116,6 +116,9 @@ def init(self, conf: ConfigTree) -> None:
         self._alchemy_extractor.init(sql_alch_conf)
         self._extract_iter: Union[None, Iterator] = None
 
+    def close(self) -> None:
+        self._alchemy_extractor.close()
+
     def extract(self) -> Union[TableMetadata, None]:
         if not self._extract_iter:
             self._extract_iter = self._get_extract_iter()

databuilder/extractor/presto_view_metadata_extractor.py

@@ -62,6 +62,9 @@ def init(self, conf: ConfigTree) -> None:
         self._alchemy_extractor.init(sql_alch_conf)
         self._extract_iter: Union[None, Iterator] = None
 
+    def close(self) -> None:
+        self._alchemy_extractor.close()
+
     def extract(self) -> Union[TableMetadata, None]:
         if not self._extract_iter:
             self._extract_iter = self._get_extract_iter()

databuilder/extractor/snowflake_metadata_extractor.py

@@ -106,6 +106,9 @@ def init(self, conf: ConfigTree) -> None:
         self._alchemy_extractor.init(sql_alch_conf)
         self._extract_iter: Union[None, Iterator] = None
 
+    def close(self) -> None:
+        self._alchemy_extractor.close()
+
     def extract(self) -> Union[TableMetadata, None]:
         if not self._extract_iter:
             self._extract_iter = self._get_extract_iter()

databuilder/extractor/snowflake_table_last_updated_extractor.py

@@ -83,6 +83,9 @@ def init(self, conf: ConfigTree) -> None:
         self._alchemy_extractor.init(sql_alch_conf)
         self._extract_iter: Union[None, Iterator] = None
 
+    def close(self) -> None:
+        self._alchemy_extractor.close()
+
     def extract(self) -> Union[TableLastUpdated, None]:
         if not self._extract_iter:
             self._extract_iter = self._get_extract_iter()

databuilder/extractor/sql_alchemy_extractor.py

@@ -38,6 +38,9 @@ def init(self, conf: ConfigTree) -> None:
             self.model_class = getattr(mod, class_name)
         self._execute_query()
 
+    def close(self) -> None:
+        self.connection.close()
+
     def _get_connection(self) -> Any:
         """
         Create a SQLAlchemy connection to Database