From e575a78893e72bc0e0c5cece0a94c041d36f0609 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Thu, 19 Mar 2026 09:56:10 +0100
Subject: [PATCH] Fix missing import, add transaction safety, clean orphaned
 snapshots

- Add load_json to tab_timeline_ng imports (NameError on disk fallback)
- Wrap save_history_tree in BEGIN/COMMIT transaction (was autocommitting
  each statement, risking partial writes on failure)
- Clean up orphaned history_snapshots in sync_to_db when nodes are
  removed from the tree

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 db.py              | 45 +++++++++++++++++++++++++++------------------
 tab_timeline_ng.py |  5 ++---
 utils.py           | 14 ++++++++++++++
 3 files changed, 43 insertions(+), 21 deletions(-)

diff --git a/db.py b/db.py
index aeb44c7..f4ac10e 100644
--- a/db.py
+++ b/db.py
@@ -326,30 +326,39 @@ class ProjectDB:
     def save_history_tree(self, data_file_id: int, tree_data: dict) -> None:
         """Save history tree, extracting node snapshots into separate table."""
         now = time.time()
-        # Extract snapshot data from nodes into history_snapshots table
         nodes = tree_data.get("nodes", {})
         slim_tree = dict(tree_data)
         slim_nodes = {}
         for nid, node in nodes.items():
-            snap = node.get("data")
-            if snap:
-                self.conn.execute(
-                    "INSERT INTO history_snapshots (data_file_id, node_id, snapshot_data, updated_at) "
-                    "VALUES (?, ?, ?, ?) "
-                    "ON CONFLICT(data_file_id, node_id) DO UPDATE SET "
-                    "snapshot_data=excluded.snapshot_data, updated_at=excluded.updated_at",
-                    (data_file_id, nid, json.dumps(snap), now),
-                )
-            # Store node without data in tree
             slim_nodes[nid] = {k: v for k, v in node.items() if k != "data"}
         slim_tree["nodes"] = slim_nodes
-        self.conn.execute(
-            "INSERT INTO history_trees (data_file_id, tree_data, updated_at) "
-            "VALUES (?, ?, ?) "
-            "ON CONFLICT(data_file_id) DO UPDATE SET tree_data=excluded.tree_data, updated_at=excluded.updated_at",
-            (data_file_id, json.dumps(slim_tree), now),
-        )
-        self.conn.commit()
+
+        self.conn.execute("BEGIN IMMEDIATE")
+        try:
+            # Extract snapshot data from nodes into history_snapshots table
+            for nid, node in nodes.items():
+                snap = node.get("data")
+                if snap:
+                    self.conn.execute(
+                        "INSERT INTO history_snapshots (data_file_id, node_id, snapshot_data, updated_at) "
+                        "VALUES (?, ?, ?, ?) "
+                        "ON CONFLICT(data_file_id, node_id) DO UPDATE SET "
+                        "snapshot_data=excluded.snapshot_data, updated_at=excluded.updated_at",
+                        (data_file_id, nid, json.dumps(snap), now),
+                    )
+            self.conn.execute(
+                "INSERT INTO history_trees (data_file_id, tree_data, updated_at) "
+                "VALUES (?, ?, ?) "
+                "ON CONFLICT(data_file_id) DO UPDATE SET tree_data=excluded.tree_data, updated_at=excluded.updated_at",
+                (data_file_id, json.dumps(slim_tree), now),
+            )
+            self.conn.execute("COMMIT")
+        except Exception:
+            try:
+                self.conn.execute("ROLLBACK")
+            except Exception:
+                pass
+            raise
 
     def get_history_tree(self, data_file_id: int) -> dict | None:
         """Load history tree metadata (without snapshot data)."""
diff --git a/tab_timeline_ng.py b/tab_timeline_ng.py
index efc4924..6dffd32 100644
--- a/tab_timeline_ng.py
+++ b/tab_timeline_ng.py
@@ -8,7 +8,7 @@ from nicegui import ui
 
 from state import AppState
 from history_tree import HistoryTree
-from utils import save_json, sync_to_db, KEY_BATCH_DATA, KEY_HISTORY_TREE
+from utils import save_json, load_json, sync_to_db, KEY_BATCH_DATA, KEY_HISTORY_TREE
 
 logger = logging.getLogger(__name__)
 
@@ -580,8 +580,7 @@ async def _restore_node(data, node, htree, file_path, state: AppState):
                 state.db.get_node_snapshot, df['id'], node['id'])
     if not raw_snap:
         # Last resort: read from JSON file on disk
-        from utils import load_json as _load_json
-        raw_file, _ = await asyncio.to_thread(_load_json, file_path)
+        raw_file, _ = await asyncio.to_thread(load_json, file_path)
         tree_on_disk = raw_file.get(KEY_HISTORY_TREE, {})
         raw_snap = tree_on_disk.get('nodes', {}).get(node['id'], {}).get('data', {})
     node_data = json.loads(json.dumps(raw_snap)) if raw_snap else {}
diff --git a/utils.py b/utils.py
index 56146f8..8b8b640 100644
--- a/utils.py
+++ b/utils.py
@@ -312,6 +312,20 @@ def sync_to_db(db, project_name: str, file_path: Path, data: dict) -> None:
                     "ON CONFLICT(data_file_id) DO UPDATE SET tree_data=excluded.tree_data, updated_at=excluded.updated_at",
                     (df_id, json.dumps(slim_tree), now),
                 )
+                # Clean up orphaned snapshots for nodes no longer in tree
+                current_node_ids = set(nodes.keys())
+                if current_node_ids:
+                    placeholders = ",".join("?" for _ in current_node_ids)
+                    db.conn.execute(
+                        f"DELETE FROM history_snapshots WHERE data_file_id = ? "
+                        f"AND node_id NOT IN ({placeholders})",
+                        (df_id, *current_node_ids),
+                    )
+                else:
+                    db.conn.execute(
+                        "DELETE FROM history_snapshots WHERE data_file_id = ?",
+                        (df_id,),
+                    )
 
             db.conn.execute("COMMIT")
         except Exception: