refactor(db): use atomic json_set for device and app field updates

IngmarStein · IngmarStein · commit 00a957d98cd8 · 2025-11-10T20:29:14.000+01:00
Previously, updating a single field on a device or app required loading the entire user JSON object, modifying it in memory, and writing the whole object back to the database. This approach was inefficient and created potential race conditions where concurrent updates could overwrite each other.

This commit refactors the update logic to use SQLite's `json_set` function for atomic and granular field updates. New helper functions, `update_device_field` and `update_app_field`, have been introduced to encapsulate this logic.

All relevant routes in the manager and the websocket receiver have been updated to use these new functions, ensuring that changes are committed atomically and safely.
diff --git a/tronbyt_server/db.py b/tronbyt_server/db.py
@@ -6,9 +6,10 @@
 import shutil
 import sqlite3
 import string
+from contextlib import contextmanager
 from datetime import date, datetime, timedelta
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, Generator
 from urllib.parse import quote, unquote
 from zoneinfo import ZoneInfo
 
@@ -26,6 +27,20 @@
 logger = logging.getLogger(__name__)
 
 
+@contextmanager
+def db_transaction(
+    db_conn: sqlite3.Connection,
+) -> Generator[sqlite3.Cursor, None, None]:
+    """A context manager for database transactions."""
+    cursor = db_conn.cursor()
+    try:
+        yield cursor
+        db_conn.commit()
+    except sqlite3.Error:
+        db_conn.rollback()
+        raise
+
+
 def get_db() -> sqlite3.Connection:
     """Get a database connection."""
     return sqlite3.connect(
@@ -971,6 +986,55 @@ def get_device_by_name(user: User, name: str) -> Device | None:
     return None
 
 
+def _update_json_field(
+    cursor: sqlite3.Cursor,
+    username: str,
+    path: str,
+    value: Any,
+) -> None:
+    """
+    Update a single field in the json_data using the provided cursor.
+    This function does NOT commit the transaction.
+    """
+    cursor.execute(
+        """
+        UPDATE json_data
+        SET data = json_set(data, ?, ?)
+        WHERE username = ?
+        """,
+        (path, value, username),
+    )
+
+
+def update_device_field(
+    cursor: sqlite3.Cursor, username: str, device_id: str, field: str, value: Any
+) -> None:
+    """
+    Update a single field for a device using the provided cursor.
+    This function does NOT commit the transaction.
+    """
+    path = f"$.devices.{device_id}.{field}"
+    _update_json_field(cursor, username, path, value)
+    logger.debug(f"Queued update for {field} for device {device_id}")
+
+
+def update_app_field(
+    cursor: sqlite3.Cursor,
+    username: str,
+    device_id: str,
+    iname: str,
+    field: str,
+    value: Any,
+) -> None:
+    """
+    Update a single field for an app using the provided cursor.
+    This function does NOT commit the transaction.
+    """
+    path = f"$.devices.{device_id}.apps.{iname}.{field}"
+    _update_json_field(cursor, username, path, value)
+    logger.debug(f"Queued update for {field} for app {iname} on device {device_id}")
+
+
 def get_device_webp_dir(device_id: str, create: bool = True) -> Path:
     """Get the WebP directory for a device."""
     path = get_data_dir() / "webp" / secure_filename(device_id)
@@ -981,19 +1045,45 @@ def get_device_webp_dir(device_id: str, create: bool = True) -> Path:
 
 def get_device_by_id(db: sqlite3.Connection, device_id: str) -> Device | None:
     """Get a device by ID."""
-    for user in get_all_users(db):
-        device = user.devices.get(device_id)
-        if device:
-            return device
+    cursor = db.cursor()
+    cursor.execute(
+        """
+        SELECT T2.value
+        FROM json_data AS T1, json_each(T1.data, '$.devices') AS T2
+        WHERE T2.key = ?
+        """,
+        (device_id,),
+    )
+    row = cursor.fetchone()
+    if row:
+        try:
+            device_data = json.loads(row[0])
+            return Device.model_validate(device_data)
+        except (json.JSONDecodeError, ValidationError) as e:
+            logger.error(f"Error processing device data for device {device_id}: {e}")
+            return None
     return None
 
 
 def get_user_by_device_id(db: sqlite3.Connection, device_id: str) -> User | None:
     """Get a user by device ID."""
-    for user in get_all_users(db):
-        device = user.devices.get(device_id)
-        if device:
-            return user
+    cursor = db.cursor()
+    cursor.execute(
+        """
+        SELECT T1.data
+        FROM json_data AS T1, json_each(T1.data, '$.devices') AS T2
+        WHERE T2.key = ?
+        """,
+        (device_id,),
+    )
+    row = cursor.fetchone()
+    if row:
+        try:
+            user_data = json.loads(row[0])
+            return User.model_validate(user_data)
+        except (json.JSONDecodeError, ValidationError) as e:
+            logger.error(f"Error processing user data for device {device_id}: {e}")
+            return None
     return None
 
 
@@ -1066,25 +1156,38 @@ def add_pushed_app(
         )
         return
 
-    device.apps[installation_id] = app
-    save_user(db, user)
+    save_app(db, device_id, app)
 
 
 def save_app(db: sqlite3.Connection, device_id: str, app: App) -> bool:
-    """Save an app."""
+    """Save an app atomically using json_set."""
+    if not app.iname:
+        return True  # Nothing to save if iname is missing
+
+    user = get_user_by_device_id(db, device_id)
+    if not user:
+        logger.error(f"Cannot save app: user not found for device {device_id}")
+        return False
+
     try:
-        user = get_user_by_device_id(db, device_id)
-        if not user:
-            return False
-        if not app.iname:
-            return True
-        device = user.devices.get(device_id)
-        if not device:
-            return False
-        device.apps[app.iname] = app
-        save_user(db, user)
+        path = f"$.devices.{device_id}.apps.{app.iname}"
+        app_json = app.model_dump_json()
+
+        cursor = db.cursor()
+        cursor.execute(
+            """
+            UPDATE json_data
+            SET data = json_set(data, ?, json(?))
+            WHERE username = ?
+            """,
+            (path, app_json, user.username),
+        )
+        db.commit()
+        logger.debug(f"Atomically saved app {app.iname} for user {user.username}")
         return True
-    except Exception:
+    except sqlite3.Error as e:
+        logger.error(f"Could not save app {app.iname} for user {user.username}: {e}")
+        db.rollback()
         return False
 
 
diff --git a/tronbyt_server/routers/manager.py b/tronbyt_server/routers/manager.py
@@ -362,8 +362,15 @@ def next_app_logic(
         # pinned app is set as the app iname in the device object, we need to clear it.
         if getattr(device, "pinned_app", None) == app.iname:
             logger.info(f"Unpinning app {app.iname} because fail or empty render")
-            device.pinned_app = ""
-            db.save_user(db_conn, user)
+            try:
+                with db.db_transaction(db_conn) as cursor:
+                    db.update_device_field(
+                        cursor, user.username, device.id, "pinned_app", ""
+                    )
+                # Keep the in-memory object in sync
+                device.pinned_app = ""
+            except sqlite3.Error as e:
+                logger.error(f"Failed to unpin app for device {device.id}: {e}")
 
         # Pass next_index directly since it already points to the next app
         return next_app_logic(db_conn, user, device, next_index, recursion_depth + 1)
@@ -376,9 +383,16 @@ def next_app_logic(
     if webp_path.exists() and webp_path.stat().st_size > 0:
         # App rendered successfully - display it and save the index
         response = send_image(webp_path, device, app)
-        # Save the next_index as the new last_app_index
-        device.last_app_index = next_index
-        db.save_user(db_conn, user)
+        # Atomically save the next_index as the new last_app_index
+        try:
+            with db.db_transaction(db_conn) as cursor:
+                db.update_device_field(
+                    cursor, user.username, device.id, "last_app_index", next_index
+                )
+            # Keep the in-memory object in sync for the current request
+            device.last_app_index = next_index
+        except sqlite3.Error as e:
+            logger.error(f"Failed to update last_app_index for device {device.id}: {e}")
         return response
 
     # WebP file doesn't exist or is empty - skip this app
@@ -578,8 +592,20 @@ async def update_brightness(
             content="Brightness must be between 0 and 5",
         )
 
-    device.brightness = db.ui_scale_to_percent(brightness)
-    db.save_user(db_conn, user)
+    percent_brightness = db.ui_scale_to_percent(brightness)
+    try:
+        with db.db_transaction(db_conn) as cursor:
+            db.update_device_field(
+                cursor, user.username, device.id, "brightness", percent_brightness
+            )
+        device.brightness = percent_brightness  # keep in-memory model updated
+    except sqlite3.Error as e:
+        logger.error(f"Failed to update brightness for device {device.id}: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Database error",
+        )
+
     new_brightness_percent = db.get_device_brightness_percent(device)
 
     # Send brightness command directly to active websocket connection (if any)
@@ -615,8 +641,18 @@ def update_interval(
     if not device:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
 
-    device.default_interval = interval
-    db.save_user(db_conn, user)
+    try:
+        with db.db_transaction(db_conn) as cursor:
+            db.update_device_field(
+                cursor, user.username, device.id, "default_interval", interval
+            )
+        device.default_interval = interval
+    except sqlite3.Error as e:
+        logger.error(f"Failed to update interval for device {device.id}: {e}")
+        return Response(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            content="Database error",
+        )
     return Response(status_code=status.HTTP_200_OK)
 
 
@@ -1031,15 +1067,22 @@ def toggle_pin(
     if iname not in device.apps:
         return Response(status_code=status.HTTP_404_NOT_FOUND, content="App not found")
 
-    # Check if this app is currently pinned
-    if getattr(device, "pinned_app", None) == iname:
-        device.pinned_app = ""
-        flash(request, _("App unpinned."))
-    else:
-        device.pinned_app = iname
-        flash(request, _("App pinned."))
+    new_pinned_app = "" if getattr(device, "pinned_app", None) == iname else iname
+
+    try:
+        with db.db_transaction(db_conn) as cursor:
+            db.update_device_field(
+                cursor, user.username, device.id, "pinned_app", new_pinned_app
+            )
+        device.pinned_app = new_pinned_app
+        if new_pinned_app:
+            flash(request, _("App pinned."))
+        else:
+            flash(request, _("App unpinned."))
+    except sqlite3.Error as e:
+        logger.error(f"Failed to toggle pin for device {device.id}: {e}")
+        flash(request, _("Error updating pin status."), "error")
 
-    db.save_user(db_conn, user)
     return RedirectResponse(url="/", status_code=status.HTTP_302_FOUND)
 
 
@@ -1296,9 +1339,23 @@ def toggle_enabled(
     if not app:
         return RedirectResponse(url="/", status_code=status.HTTP_404_NOT_FOUND)
 
-    app.enabled = not app.enabled
-    db.save_user(db_conn, user)
-    flash(request, _("Changes saved."))
+    new_enabled_status = not app.enabled
+    try:
+        with db.db_transaction(db_conn) as cursor:
+            db.update_app_field(
+                cursor,
+                user.username,
+                device.id,
+                app.iname,
+                "enabled",
+                new_enabled_status,
+            )
+        app.enabled = new_enabled_status
+        flash(request, _("Changes saved."))
+    except sqlite3.Error as e:
+        logger.error(f"Failed to toggle enabled for app {app.iname}: {e}")
+        flash(request, _("Error saving changes."), "error")
+
     return RedirectResponse(url="/", status_code=status.HTTP_302_FOUND)
 
 
@@ -2094,9 +2151,30 @@ def next_app(
     user = deps.user
     device = deps.device
 
-    device.last_seen = datetime.now(timezone.utc)
-    device.info.protocol_type = ProtocolType.HTTP
-    db.save_user(db_conn, user)
+    now = datetime.now(timezone.utc)
+    try:
+        with db.db_transaction(db_conn) as cursor:
+            db.update_device_field(
+                cursor,
+                user.username,
+                device.id,
+                "last_seen",
+                now.isoformat(),
+            )
+            db.update_device_field(
+                cursor,
+                user.username,
+                device.id,
+                "info.protocol_type",
+                ProtocolType.HTTP.value,
+            )
+        # Keep in-memory object in sync
+        device.last_seen = now
+        device.info.protocol_type = ProtocolType.HTTP
+    except sqlite3.Error as e:
+        logger.error(
+            f"Failed to update device {device.id} last_seen and protocol_type: {e}"
+        )
 
     return next_app_logic(db_conn, user, device)
 
diff --git a/tronbyt_server/routers/websockets.py b/tronbyt_server/routers/websockets.py
