diff --git a/databuilder/extractor/athena_metadata_extractor.py b/databuilder/extractor/athena_metadata_extractor.py index 9ef75e31a..83785b044 100644 --- a/databuilder/extractor/athena_metadata_extractor.py +++ b/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() diff --git a/databuilder/extractor/druid_metadata_extractor.py b/databuilder/extractor/druid_metadata_extractor.py index b2f7766ba..2d918cbb6 100644 --- a/databuilder/extractor/druid_metadata_extractor.py +++ b/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() diff --git a/databuilder/extractor/mssql_metadata_extractor.py b/databuilder/extractor/mssql_metadata_extractor.py index ce89752a0..e19bf0eb1 100644 --- a/databuilder/extractor/mssql_metadata_extractor.py +++ b/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() diff --git a/databuilder/extractor/presto_view_metadata_extractor.py b/databuilder/extractor/presto_view_metadata_extractor.py index 0a3452774..e33bfe848 100644 --- a/databuilder/extractor/presto_view_metadata_extractor.py +++ b/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() diff --git a/databuilder/extractor/snowflake_metadata_extractor.py b/databuilder/extractor/snowflake_metadata_extractor.py index 33f0eda3b..df1b2a414 100644 --- a/databuilder/extractor/snowflake_metadata_extractor.py +++ b/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() diff --git a/databuilder/extractor/snowflake_table_last_updated_extractor.py b/databuilder/extractor/snowflake_table_last_updated_extractor.py index d477cece4..627095dca 100644 --- a/databuilder/extractor/snowflake_table_last_updated_extractor.py +++ b/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() diff --git a/databuilder/extractor/sql_alchemy_extractor.py b/databuilder/extractor/sql_alchemy_extractor.py index a375781a5..a8b64a134 100644 --- a/databuilder/extractor/sql_alchemy_extractor.py +++ b/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