feat(app): expose a cursor, not a connection in db util

Author: psychedelicious

invokeai/app/services/board_image_records/board_image_records_sqlite.py

@@ -21,8 +21,8 @@ def add_image_to_board(
         board_id: str,
         image_name: str,
     ) -> None:
-        with self._db.conn() as conn:
-            conn.execute(
+        with self._db.transaction() as cursor:
+            cursor.execute(
                 """--sql
                 INSERT INTO board_images (board_id, image_name)
                 VALUES (?, ?)
@@ -35,8 +35,8 @@ def remove_image_from_board(
         self,
         image_name: str,
     ) -> None:
-        with self._db.conn() as conn:
-            conn.execute(
+        with self._db.transaction() as cursor:
+            cursor.execute(
                 """--sql
                 DELETE FROM board_images
                 WHERE image_name = ?;
@@ -50,8 +50,7 @@ def get_images_for_board(
         offset: int = 0,
         limit: int = 10,
     ) -> OffsetPaginatedResults[ImageRecord]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT images.*
@@ -80,8 +79,7 @@ def get_all_board_image_names_for_board(
         categories: list[ImageCategory] | None,
         is_intermediate: bool | None,
     ) -> list[str]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             params: list[str | bool] = []
 
             # Base query is a join between images and board_images
@@ -137,8 +135,7 @@ def get_board_for_image(
         self,
         image_name: str,
     ) -> Optional[str]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                     SELECT board_id
@@ -153,8 +150,7 @@ def get_board_for_image(
         return cast(str, result[0])
 
     def get_image_count_for_board(self, board_id: str) -> int:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                     SELECT COUNT(*)

invokeai/app/services/board_records/board_records_sqlite.py

@@ -23,9 +23,8 @@ def __init__(self, db: SqliteDatabase) -> None:
         self._db = db
 
     def delete(self, board_id: str) -> None:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     DELETE FROM boards
@@ -40,10 +39,9 @@ def save(
         self,
         board_name: str,
     ) -> BoardRecord:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
                 board_id = uuid_string()
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     INSERT OR IGNORE INTO boards (board_id, board_name)
@@ -59,9 +57,8 @@ def get(
         self,
         board_id: str,
     ) -> BoardRecord:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     SELECT *
@@ -83,9 +80,8 @@ def update(
         board_id: str,
         changes: BoardChanges,
     ) -> BoardRecord:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 # Change the name of a board
                 if changes.board_name is not None:
                     cursor.execute(
@@ -131,9 +127,7 @@ def get_many(
         limit: int = 10,
         include_archived: bool = False,
     ) -> OffsetPaginatedResults[BoardRecord]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
-
+        with self._db.transaction() as cursor:
             # Build base query
             base_query = """
                     SELECT *
@@ -179,8 +173,7 @@ def get_many(
     def get_all(
         self, order_by: BoardRecordOrderBy, direction: SQLiteDirection, include_archived: bool = False
     ) -> list[BoardRecord]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             if order_by == BoardRecordOrderBy.Name:
                 base_query = """
                         SELECT *

invokeai/app/services/image_records/image_records_sqlite.py

@@ -27,9 +27,8 @@ def __init__(self, db: SqliteDatabase) -> None:
         self._db = db
 
     def get(self, image_name: str) -> ImageRecord:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     f"""--sql
                     SELECT {IMAGE_DTO_COLS} FROM images
@@ -48,9 +47,8 @@ def get(self, image_name: str) -> ImageRecord:
         return deserialize_image_record(dict(result))
 
     def get_metadata(self, image_name: str) -> Optional[MetadataField]:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     SELECT metadata FROM images
@@ -76,9 +74,8 @@ def update(
         image_name: str,
         changes: ImageRecordChanges,
     ) -> None:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 # Change the category of the image
                 if changes.image_category is not None:
                     cursor.execute(
@@ -138,9 +135,7 @@ def get_many(
         board_id: Optional[str] = None,
         search_term: Optional[str] = None,
     ) -> OffsetPaginatedResults[ImageRecord]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
-
+        with self._db.transaction() as cursor:
             # Manually build two queries - one for the count, one for the records
             count_query = """--sql
             SELECT COUNT(*)
@@ -227,20 +222,20 @@ def get_many(
             # Build the list of images, deserializing each row
             cursor.execute(images_query, images_params)
             result = cast(list[sqlite3.Row], cursor.fetchall())
-        images = [deserialize_image_record(dict(r)) for r in result]
 
-        # Set up and execute the count query, without pagination
-        count_query += query_conditions + ";"
-        count_params = query_params.copy()
-        cursor.execute(count_query, count_params)
-        count = cast(int, cursor.fetchone()[0])
+            images = [deserialize_image_record(dict(r)) for r in result]
+
+            # Set up and execute the count query, without pagination
+            count_query += query_conditions + ";"
+            count_params = query_params.copy()
+            cursor.execute(count_query, count_params)
+            count = cast(int, cursor.fetchone()[0])
 
         return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count)
 
     def delete(self, image_name: str) -> None:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     DELETE FROM images
@@ -252,10 +247,8 @@ def delete(self, image_name: str) -> None:
                 raise ImageRecordDeleteException from e
 
     def delete_many(self, image_names: list[str]) -> None:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
-
                 placeholders = ",".join("?" for _ in image_names)
 
                 # Construct the SQLite query with the placeholders
@@ -268,8 +261,7 @@ def delete_many(self, image_names: list[str]) -> None:
                 raise ImageRecordDeleteException from e
 
     def get_intermediates_count(self) -> int:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT COUNT(*) FROM images
@@ -280,9 +272,8 @@ def get_intermediates_count(self) -> int:
         return count
 
     def delete_intermediates(self) -> list[str]:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     SELECT image_name FROM images
@@ -315,9 +306,8 @@ def save(
         node_id: Optional[str] = None,
         metadata: Optional[str] = None,
     ) -> datetime:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     INSERT OR IGNORE INTO images (
@@ -366,8 +356,7 @@ def save(
         return created_at
 
     def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT images.*
@@ -398,9 +387,7 @@ def get_image_names(
         board_id: Optional[str] = None,
         search_term: Optional[str] = None,
     ) -> ImageNamesResult:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
-
+        with self._db.transaction() as cursor:
             # Build query conditions (reused for both starred count and image names queries)
             query_conditions = ""
             query_params: list[Union[int, str, bool]] = []

invokeai/app/services/model_records/model_records_sql.py

@@ -88,9 +88,8 @@ def add_model(self, config: AnyModelConfig) -> AnyModelConfig:
 
         Can raise DuplicateModelException and InvalidModelConfigException exceptions.
         """
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             try:
-                cursor = conn.cursor()
                 cursor.execute(
                     """--sql
                     INSERT INTO models (
@@ -127,8 +126,7 @@ def del_model(self, key: str) -> None:
 
         Can raise an UnknownModelException
         """
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 DELETE FROM models
@@ -140,7 +138,7 @@ def del_model(self, key: str) -> None:
                 raise UnknownModelException("model not found")
 
     def update_model(self, key: str, changes: ModelRecordChanges) -> AnyModelConfig:
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             record = self.get_model(key)
 
             # Model configs use pydantic's `validate_assignment`, so each change is validated by pydantic.
@@ -149,7 +147,6 @@ def update_model(self, key: str, changes: ModelRecordChanges) -> AnyModelConfig:
 
             json_serialized = record.model_dump_json()
 
-            cursor = conn.cursor()
             cursor.execute(
                 """--sql
                 UPDATE models
@@ -172,8 +169,7 @@ def get_model(self, key: str) -> AnyModelConfig:
 
         Exceptions: UnknownModelException
         """
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT config, strftime('%s',updated_at) FROM models
@@ -188,8 +184,7 @@ def get_model(self, key: str) -> AnyModelConfig:
         return model
 
     def get_model_by_hash(self, hash: str) -> AnyModelConfig:
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT config, strftime('%s',updated_at) FROM models
@@ -209,8 +204,7 @@ def exists(self, key: str) -> bool:
 
         :param key: Unique key for the model to be deleted
         """
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 select count(*) FROM models
@@ -241,7 +235,7 @@ def search_by_attr(
         If none of the optional filters are passed, will return all
         models in the database.
         """
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             assert isinstance(order_by, ModelRecordOrderBy)
             ordering = {
                 ModelRecordOrderBy.Default: "type, base, name, format",
@@ -267,7 +261,6 @@ def search_by_attr(
                 bindings.append(model_format)
             where = f"WHERE {' AND '.join(where_clause)}" if where_clause else ""
 
-            cursor = conn.cursor()
             cursor.execute(
                 f"""--sql
                 SELECT config, strftime('%s',updated_at)
@@ -299,8 +292,7 @@ def search_by_attr(
 
     def search_by_path(self, path: Union[str, Path]) -> List[AnyModelConfig]:
         """Return models with the indicated path."""
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT config, strftime('%s',updated_at) FROM models
@@ -313,8 +305,7 @@ def search_by_path(self, path: Union[str, Path]) -> List[AnyModelConfig]:
 
     def search_by_hash(self, hash: str) -> List[AnyModelConfig]:
         """Return models with the indicated hash."""
-        with self._db.conn() as conn:
-            cursor = conn.cursor()
+        with self._db.transaction() as cursor:
             cursor.execute(
                 """--sql
                 SELECT config, strftime('%s',updated_at) FROM models
@@ -329,7 +320,7 @@ def list_models(
         self, page: int = 0, per_page: int = 10, order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default
     ) -> PaginatedResults[ModelSummary]:
         """Return a paginated summary listing of each model in the database."""
-        with self._db.conn() as conn:
+        with self._db.transaction() as cursor:
             assert isinstance(order_by, ModelRecordOrderBy)
             ordering = {
                 ModelRecordOrderBy.Default: "type, base, name, format",
@@ -339,8 +330,6 @@ def list_models(
                 ModelRecordOrderBy.Format: "format",
             }
 
-            cursor = conn.cursor()
-
             # Lock so that the database isn't updated while we're doing the two queries.
             # query1: get the total number of model configs
             cursor.execute(