From 39f07ec2501d0b5b22c6a2bb3c62602d168b578d Mon Sep 17 00:00:00 2001
From: Logan Ward <WardLT@users.noreply.github.com>
Date: Fri, 1 Dec 2023 20:16:59 -0500
Subject: [PATCH] Move logging, shutdown tpex

---
 examol/store/db/base.py   | 1 -
 examol/store/db/memory.py | 5 ++++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/examol/store/db/base.py b/examol/store/db/base.py
index 750d916..5c797ae 100644
--- a/examol/store/db/base.py
+++ b/examol/store/db/base.py
@@ -80,7 +80,6 @@ def export_records(self, path: Path):
         Args:
             path: Path in which to save all data. Use a ".json.gz"
         """
-        logger.info(f'Started writing to {path}')
         with (gzip.open(path, 'wt') if path.name.endswith('.gz') else open(path, 'w')) as fp:
             for record in self.iterate_over_records():
                 print(record.json(), file=fp)
diff --git a/examol/store/db/memory.py b/examol/store/db/memory.py
index 0d30805..22d6a5b 100644
--- a/examol/store/db/memory.py
+++ b/examol/store/db/memory.py
@@ -32,7 +32,7 @@ def __init__(self, path: Path | None, write_freq: float = 10.):
         self.db: dict[str, MoleculeRecord] = {}
 
         # Start thread which writes until
-        self._thread_pool = ThreadPoolExecutor(max_workers=1)
+        self._thread_pool ThreadPoolExecutor | None = None
         self._write_thread: Future | None = None
         self._updates_available: Event = Event()
         self._closing = Event()
@@ -43,6 +43,7 @@ def __init__(self, path: Path | None, write_freq: float = 10.):
     def __enter__(self):
         if self.path is not None:
             logger.info('Start the writing thread')
+            self._thread_pool = ThreadPoolExecutor(max_workers=1)
             self._write_thread = self._thread_pool.submit(self._writer)
 
             # Add a callback to print a logging message if there is an error
@@ -64,6 +65,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # Mark that we're closed
             self._write_thread = None
             self._closing.clear()
+            self._thread_pool.shutdown()
 
     def _load_molecules(self):
         """Load molecules from disk"""
@@ -105,6 +107,7 @@ def _writer(self):
 
             # Checkpoint and advance the standoff
             temp_path = self.path.parent / (self.path.name + "-new")
+            logger.info(f'Started writing to {temp_path}')
             self.export_records(temp_path)
             move(temp_path, self.path)
             next_write = monotonic() + self.write_freq