From 0228aba06f38321456aa57f07e554629712d785c Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 30 Nov 2023 20:07:53 +1100
Subject: [PATCH] feat(backend): display freed space when cleaning DB

---
 .../latents_storage/latents_storage_disk.py   |  2 +-
 invokeai/app/services/shared/sqlite.py        | 21 ++++++++++---------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/invokeai/app/services/latents_storage/latents_storage_disk.py b/invokeai/app/services/latents_storage/latents_storage_disk.py
index 58d1df5081..f59292bb91 100644
--- a/invokeai/app/services/latents_storage/latents_storage_disk.py
+++ b/invokeai/app/services/latents_storage/latents_storage_disk.py
@@ -33,7 +33,7 @@ class DiskLatentsStorage(LatentsStorageBase):
         if deleted_latents_count > 0:
             freed_space_in_mb = round(freed_space / 1024 / 1024, 2)
             self._invoker.services.logger.info(
-                f"Deleted {deleted_latents_count} latents files, freeing {freed_space_in_mb}MB"
+                f"Deleted {deleted_latents_count} latents files (freed {freed_space_in_mb}MB)"
             )
 
     def get(self, name: str) -> torch.Tensor:
diff --git a/invokeai/app/services/shared/sqlite.py b/invokeai/app/services/shared/sqlite.py
index 3c75c3d6a7..511023bd8a 100644
--- a/invokeai/app/services/shared/sqlite.py
+++ b/invokeai/app/services/shared/sqlite.py
@@ -1,6 +1,7 @@
 import sqlite3
 import threading
 from logging import Logger
+from pathlib import Path
 
 from invokeai.app.services.config import InvokeAIAppConfig
 
@@ -8,25 +9,20 @@ sqlite_memory = ":memory:"
 
 
 class SqliteDatabase:
-    conn: sqlite3.Connection
-    lock: threading.RLock
-    _logger: Logger
-    _config: InvokeAIAppConfig
-
     def __init__(self, config: InvokeAIAppConfig, logger: Logger):
         self._logger = logger
         self._config = config
 
         if self._config.use_memory_db:
-            location = sqlite_memory
+            self.db_path = sqlite_memory
             logger.info("Using in-memory database")
         else:
             db_path = self._config.db_path
             db_path.parent.mkdir(parents=True, exist_ok=True)
-            location = str(db_path)
-            self._logger.info(f"Using database at {location}")
+            self.db_path = str(db_path)
+            self._logger.info(f"Using database at {self.db_path}")
 
-        self.conn = sqlite3.connect(location, check_same_thread=False)
+        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
         self.lock = threading.RLock()
         self.conn.row_factory = sqlite3.Row
 
@@ -37,10 +33,15 @@ class SqliteDatabase:
 
     def clean(self) -> None:
         try:
+            if self.db_path == sqlite_memory:
+                return
+            initial_db_size = Path(self.db_path).stat().st_size
             self.lock.acquire()
             self.conn.execute("VACUUM;")
             self.conn.commit()
-            self._logger.info("Cleaned database")
+            final_db_size = Path(self.db_path).stat().st_size
+            freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2)
+            self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
         except Exception as e:
             self._logger.error(f"Error cleaning database: {e}")
             raise e