feat(db): decouple SqliteDatabase from config object

- Simplify init args to path (None means use memory), logger, and verbose
- Add docstrings to SqliteDatabase (it had almost none)
- Update all usages of the class

invokeai/app/api/dependencies.py

@@ -2,7 +2,6 @@
 
 from functools import partial
 from logging import Logger
-from pathlib import Path
 
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_1 import migration_1
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import migration_2
@@ -75,10 +74,11 @@ class ApiDependencies:
         output_folder = config.output_path
         image_files = DiskImageFileStorage(f"{output_folder}/images")
 
-        db = SqliteDatabase(config, logger)
+        db_path = None if config.use_memory_db else config.db_path
+        db = SqliteDatabase(db_path=db_path, logger=logger, verbose=config.log_sql)
 
         migrator = SQLiteMigrator(
-            db_path=db.database if isinstance(db.database, Path) else None,
+            db_path=db.db_path,
             conn=db.conn,
             lock=db.lock,
             logger=logger,
@@ -89,7 +89,10 @@ class ApiDependencies:
         migrator.register_migration(migration_2)
         did_migrate = migrator.run_migrations()
 
-        if not db.is_memory and did_migrate:
+        # We need to reinitialize the database if we migrated, but only if we are using a file database.
+        # This closes the SqliteDatabase's connection and re-runs its `__init__` logic.
+        # If we do this with a memory database, we wipe the db.
+        if not db.db_path and did_migrate:
             db.reinitialize()
 
         configuration = config

invokeai/app/services/shared/sqlite/sqlite_database.py

@@ -3,58 +3,83 @@ import threading
 from logging import Logger
 from pathlib import Path
 
-from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.shared.sqlite.sqlite_common import sqlite_memory
 
 
 class SqliteDatabase:
-    database: Path | str  # Must declare this here to satisfy type checker
+    """
+    Manages a connection to an SQLite database.
 
-    def __init__(self, config: InvokeAIAppConfig, logger: Logger) -> None:
+    This is a light wrapper around the `sqlite3` module, providing a few conveniences:
-        self.initialize(config, logger)
+    - The database file is written to disk if it does not exist.
+    - Foreign key constraints are enabled by default.
+    - The connection is configured to use the `sqlite3.Row` row factory.
+    - A `conn` attribute is provided to access the connection.
+    - A `lock` attribute is provided to lock the database connection.
+    - A `clean` method to run the VACUUM command and report on the freed space.
+    - A `reinitialize` method to close the connection and re-run the init.
+    - A `close` method to close the connection.
 
-    def initialize(self, config: InvokeAIAppConfig, logger: Logger) -> None:
+    :param db_path: Path to the database file. If None, an in-memory database is used.
-        self._logger = logger
+    :param logger: Logger to use for logging.
-        self._config = config
+    :param verbose: Whether to log SQL statements. Provides `logger.debug` as the SQLite trace callback.
-        self.is_memory = False
+    """
-        if self._config.use_memory_db:
+
-            self.database = sqlite_memory
+    def __init__(self, db_path: Path | None, logger: Logger, verbose: bool = False) -> None:
-            self.is_memory = True
+        self.initialize(db_path=db_path, logger=logger, verbose=verbose)
+
+    def initialize(self, db_path: Path | None, logger: Logger, verbose: bool = False) -> None:
+        """Initializes the database. This is used internally by the class constructor."""
+        self.logger = logger
+        self.db_path = db_path
+        self.verbose = verbose
+
+        if not self.db_path:
             logger.info("Initializing in-memory database")
         else:
-            self.database = self._config.db_path
+            self.db_path.parent.mkdir(parents=True, exist_ok=True)
-            self.database.parent.mkdir(parents=True, exist_ok=True)
+            self.logger.info(f"Initializing database at {self.db_path}")
-            self._logger.info(f"Initializing database at {self.database}")
 
-        self.conn = sqlite3.connect(database=self.database, check_same_thread=False)
+        self.conn = sqlite3.connect(database=self.db_path or sqlite_memory, check_same_thread=False)
         self.lock = threading.RLock()
         self.conn.row_factory = sqlite3.Row
 
-        if self._config.log_sql:
+        if self.verbose:
-            self.conn.set_trace_callback(self._logger.debug)
+            self.conn.set_trace_callback(self.logger.debug)
 
         self.conn.execute("PRAGMA foreign_keys = ON;")
 
     def reinitialize(self) -> None:
-        """Reinitializes the database. Needed after migration."""
+        """
+        Re-initializes the database by closing the connection and re-running the init.
+        Warning: This will wipe the database if it is an in-memory database.
+        """
         self.close()
-        self.initialize(self._config, self._logger)
+        self.initialize(db_path=self.db_path, logger=self.logger, verbose=self.verbose)
 
     def close(self) -> None:
+        """
+        Closes the connection to the database.
+        Warning: This will wipe the database if it is an in-memory database.
+        """
         self.conn.close()
 
     def clean(self) -> None:
+        """
+        Cleans the database by running the VACUUM command, reporting on the freed space.
+        """
+        # No need to clean in-memory database
+        if not self.db_path:
+            return
         with self.lock:
             try:
-                if self.database == sqlite_memory:
+                initial_db_size = Path(self.db_path).stat().st_size
-                    return
-                initial_db_size = Path(self.database).stat().st_size
                 self.conn.execute("VACUUM;")
                 self.conn.commit()
-                final_db_size = Path(self.database).stat().st_size
+                final_db_size = Path(self.db_path).stat().st_size
                 freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2)
                 if freed_space_in_mb > 0:
-                    self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
+                    self.logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
             except Exception as e:
-                self._logger.error(f"Error cleaning database: {e}")
+                self.logger.error(f"Error cleaning database: {e}")
                 raise

invokeai/backend/model_manager/migrate_to_db.py

@@ -47,7 +47,8 @@ class MigrateModelYamlToDb:
 
     def get_db(self) -> ModelRecordServiceSQL:
         """Fetch the sqlite3 database for this installation."""
-        db = SqliteDatabase(self.config, self.logger)
+        db_path = None if self.config.use_memory_db else self.config.db_path
+        db = SqliteDatabase(db_path=db_path, logger=self.logger, verbose=self.config.log_sql)
         return ModelRecordServiceSQL(db)
 
     def get_yaml(self) -> DictConfig:

tests/aa_nodes/test_graph_execution_state.py

@@ -1,5 +1,4 @@
 import logging
-from pathlib import Path
 
 import pytest
 
@@ -54,9 +53,10 @@ def simple_graph():
 def mock_services() -> InvocationServices:
     configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0)
     logger = InvokeAILogger.get_logger()
-    db = SqliteDatabase(configuration, logger)
+    db_path = None if configuration.use_memory_db else configuration.db_path
+    db = SqliteDatabase(db_path=db_path, logger=logger, verbose=configuration.log_sql)
     migrator = SQLiteMigrator(
-        db_path=db.database if isinstance(db.database, Path) else None,
+        db_path=db.db_path,
         conn=db.conn,
         lock=db.lock,
         logger=logger,

tests/aa_nodes/test_invoker.py

@@ -1,5 +1,4 @@
 import logging
-from pathlib import Path
 
 import pytest
 
@@ -58,9 +57,10 @@ def graph_with_subgraph():
 def mock_services() -> InvocationServices:
     configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0)
     logger = InvokeAILogger.get_logger()
-    db = SqliteDatabase(configuration, logger)
+    db_path = None if configuration.use_memory_db else configuration.db_path
+    db = SqliteDatabase(db_path=db_path, logger=logger, verbose=configuration.log_sql)
     migrator = SQLiteMigrator(
-        db_path=db.database if isinstance(db.database, Path) else None,
+        db_path=db.db_path,
         conn=db.conn,
         lock=db.lock,
         logger=logger,

tests/aa_nodes/test_sqlite.py

@@ -15,8 +15,11 @@ class TestModel(BaseModel):
 
 @pytest.fixture
 def db() -> SqliteItemStorage[TestModel]:
-    sqlite_db = SqliteDatabase(InvokeAIAppConfig(use_memory_db=True), InvokeAILogger.get_logger())
+    config = InvokeAIAppConfig(use_memory_db=True)
-    sqlite_item_storage = SqliteItemStorage[TestModel](db=sqlite_db, table_name="test", id_field="id")
+    logger = InvokeAILogger.get_logger()
+    db_path = None if config.use_memory_db else config.db_path
+    db = SqliteDatabase(db_path=db_path, logger=logger, verbose=config.log_sql)
+    sqlite_item_storage = SqliteItemStorage[TestModel](db=db, table_name="test", id_field="id")
     return sqlite_item_storage
 
 

tests/app/services/model_install/test_model_install.py

@@ -42,11 +42,12 @@ def app_config(datadir: Path) -> InvokeAIAppConfig:
 @pytest.fixture
 def store(app_config: InvokeAIAppConfig) -> ModelRecordServiceBase:
     logger = InvokeAILogger.get_logger(config=app_config)
-    database = SqliteDatabase(app_config, logger)
+    db_path = None if app_config.use_memory_db else app_config.db_path
+    db = SqliteDatabase(db_path=db_path, logger=logger, verbose=app_config.log_sql)
     migrator = SQLiteMigrator(
-        db_path=database.database if isinstance(database.database, Path) else None,
+        db_path=db.db_path,
-        conn=database.conn,
+        conn=db.conn,
-        lock=database.lock,
+        lock=db.lock,
         logger=logger,
         log_sql=app_config.log_sql,
     )
@@ -54,8 +55,8 @@ def store(app_config: InvokeAIAppConfig) -> ModelRecordServiceBase:
     migrator.register_migration(migration_2)
     migrator.run_migrations()
     # this test uses a file database, so we need to reinitialize it after migrations
-    database.reinitialize()
+    db.reinitialize()
-    store: ModelRecordServiceBase = ModelRecordServiceSQL(database)
+    store: ModelRecordServiceBase = ModelRecordServiceSQL(db)
     return store
 
 

tests/app/services/model_records/test_model_records_sql.py

@@ -3,7 +3,6 @@ Test the refactored model config classes.
 """
 
 from hashlib import sha256
-from pathlib import Path
 from typing import Any
 
 import pytest
@@ -34,9 +33,10 @@ from invokeai.backend.util.logging import InvokeAILogger
 def store(datadir: Any) -> ModelRecordServiceBase:
     config = InvokeAIAppConfig(root=datadir)
     logger = InvokeAILogger.get_logger(config=config)
-    db = SqliteDatabase(config, logger)
+    db_path = None if config.use_memory_db else config.db_path
+    db = SqliteDatabase(db_path=db_path, logger=logger, verbose=config.log_sql)
     migrator = SQLiteMigrator(
-        db_path=db.database if isinstance(db.database, Path) else None,
+        db_path=db.db_path,
         conn=db.conn,
         lock=db.lock,
         logger=logger,