From 39ef7556703da650ed1974a016df22a7b5baced1 Mon Sep 17 00:00:00 2001 From: Bee Webb Date: Wed, 28 Aug 2024 10:44:55 +0000 Subject: [PATCH] Close connections to db, and print any exceptions encountered Were were running out of db connections, but not seeing the exception message, because it's only raised in the original process if we ask for the result. --- iatidata/__init__.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/iatidata/__init__.py b/iatidata/__init__.py index b39adc9..a0e59f1 100644 --- a/iatidata/__init__.py +++ b/iatidata/__init__.py @@ -11,6 +11,7 @@ import subprocess import tempfile import time +import traceback import zipfile from collections import defaultdict from datetime import datetime @@ -305,7 +306,9 @@ def load_dataset(dataset: iatikit.Dataset) -> None: path = pathlib.Path(dataset.data_path) prefix, filename = path.parts[-2:] - with get_engine().begin() as connection: + engine = get_engine() + + with engine.begin() as connection: connection.execute( insert( table( @@ -334,6 +337,8 @@ def load_dataset(dataset: iatikit.Dataset) -> None: ) ) + engine.dispose() + def create_database_schema(): if schema: @@ -363,6 +368,12 @@ def load(processes: int, sample: Optional[int] = None) -> None: executor.submit(load_dataset, dataset) for dataset in datasets_sample ] concurrent.futures.wait(futures) + # We have to get the result (even though we don't use it) in order to get the exceptions + for future in futures: + try: + future.result() + except Exception: + print(traceback.format_exc()) engine = get_engine() with engine.begin() as connection: