Merge pull request #816 from Sameer6305/fix/782-track-entity-atomic-write

fix(provenance): make track_entity's two-step write atomic (closes #782)

CHANGELOG.md

@@ -43,6 +43,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- **`ProvenanceManager.track_entity` persisted partial history and returned fabricated entries on storage failure** (#782, #816) by @Sameer6305 and @KaifAhmad1
+  - `track_entity()`'s two-step write (history archive + primary update) is now atomic — if either write fails, the whole operation rolls back via the existing #807 `transaction()` mechanism, instead of silently persisting a partial state
+  - `track_entity()`'s return type is now `Optional[ProvenanceEntry]`: on failure it returns a safe deep copy of the pre-failure existing entry (if one existed) or `None` (if this was a brand-new, never-successfully-tracked entity) — never a fabricated object claiming values that were never actually persisted
+  - This is a behavior change for callers that inspect the return value without checking for `None` first — audited: 0 of 47 production call sites in the repo currently dereference the return value, so this is safe today, but any NEW caller must handle `None`
+  - `InMemoryStorage` gained real transactional rollback (staging-buffer based) to match this guarantee — previously `transaction()` was a no-op
+
 - **`ProvenanceManager` duplicated the same checksum/persist/exception-swallow block across 4 tracking methods** (#784, #815) by @Sameer6305 and @KaifAhmad1
   - Consolidated the repeated `entry.checksum = compute_checksum(entry)` / `try: self.storage.store(entry) except Exception: pass` block used by `track_entity`, `track_relationship`, `track_chunk`, and `track_property_source` into a single `ProvenanceManager._save_entry()` helper, preserving the existing graceful-failure behavior and the batch `_conn`/re-raise semantics from #807
   - Added 4 regression tests (`tests/provenance/test_manager.py`) covering storage-failure swallowing for each of the four tracking methods, none of which had coverage for this path before

docs/guides/provenance.md

@@ -89,7 +89,7 @@ For any regulated deployment — security operations, clinical data, financial r
 
 ## Recording provenance when ingesting data
 
-The moment data enters your graph is the moment provenance must be recorded. `track_entity()` captures the source document, the timestamp, the operator or pipeline that ran the extraction, a verbatim quote from the source, and a confidence score. It returns a `ProvenanceEntry` with a SHA-256 checksum computed automatically.
+The moment data enters your graph is the moment provenance must be recorded. `track_entity()` captures the source document, the timestamp, the operator or pipeline that ran the extraction, a verbatim quote from the source, and a confidence score. It returns an `Optional[ProvenanceEntry]` (`ProvenanceEntry` on success, or `None` if storage fails on a brand-new entity) with a SHA-256 checksum computed automatically.
 
 ```python
 # Ingesting CVE-2024-3400 from NVD and a commercial feed

docs/reference/provenance.md

@@ -221,7 +221,7 @@ cleared = manager.clear()
 
 | Method | Returns | Description |
 | :------ | :------- | :----------- |
-| `track_entity(entity_id, source, metadata, **kwargs)` | `ProvenanceEntry` | Record entity provenance; checksum set automatically |
+| `track_entity(entity_id, source, metadata, **kwargs)` | `Optional[ProvenanceEntry]` | Record entity provenance atomically; returns `ProvenanceEntry` on success, or `None`/existing entry on storage failure |
 | `track_relationship(relationship_id, source, metadata, **kwargs)` | `ProvenanceEntry` | Record relationship provenance |
 | `track_chunk(chunk_id, source_document, ...)` | `ProvenanceEntry` | Record chunk provenance with char offsets |
 | `track_property_source(entity_id, property_name, value, source)` | `ProvenanceEntry` | Record property-level source attribution |

semantica/provenance/manager.py

@@ -36,6 +36,7 @@ import threading
 from .schemas import ProvenanceEntry, SourceReference
 from .storage import ProvenanceStorage, InMemoryStorage, SQLiteStorage
 from .integrity import compute_checksum, verify_checksum
+from ..utils.logging import get_logger
 
 
 @contextmanager
@@ -112,6 +113,7 @@ class ProvenanceManager:
             storage_path: Path to SQLite database (optional, uses in-memory if None)
             config: Configuration dictionary or mapping (optional)
         """
+        self.logger = get_logger("provenance_manager")
         if storage:
             self.storage = storage
             return
@@ -179,7 +181,7 @@ class ProvenanceManager:
         metadata: Optional[Dict[str, Any]] = None,
         _conn: Optional[Any] = None,
         **kwargs
-    ) -> ProvenanceEntry:
+    ) -> Optional[ProvenanceEntry]:
         """
         Track entity provenance (kg.ProvenanceTracker compatible).
         
@@ -190,7 +192,8 @@ class ProvenanceManager:
             **kwargs: Additional fields (confidence, source_location, etc.)
             
         Returns:
-            ProvenanceEntry object
+            Optional[ProvenanceEntry]: ProvenanceEntry on success, deepcopy of existing
+            entry if update fails, or None if brand-new entity storage fails
             
         Example:
             >>> prov_mgr.track_entity(
@@ -210,6 +213,7 @@ class ProvenanceManager:
         # retrieve and store operations share a single transaction. With BEGIN IMMEDIATE,
         # concurrent calls serialize during retrieval so no intervening versions are lost.
         # If any step raises, the whole operation rolls back.
+        existing = None
         entry = None
         try:
             with self._get_or_create_transaction(_conn) as conn:
@@ -266,29 +270,24 @@ class ProvenanceManager:
                 if archived_history_id and explicit_parent_supplied:
                     entry.used_entities.append(archived_history_id)
                 
-                self._save_entry(entry, _conn=conn, _raise_on_error=(_conn is not None))
-        except Exception:
+                self._save_entry(entry, _conn=conn, _raise_on_error=True)
+        except Exception as e:
             # When called from a batch's shared transaction (_conn is not None),
             # propagate so the caller's per-item try/except can skip counting
             # this item instead of reporting an unpersisted entry as tracked (#807).
             if _conn is not None:
                 raise
-            if entry is None:
-                entry = ProvenanceEntry(
-                    entity_id=entity_id,
-                    entity_type=kwargs.get("entity_type", "entity"),
-                    activity_id=kwargs.get("activity_id", "entity_tracking"),
-                    source_document=source,
-                    source_location=kwargs.get("source_location"),
-                    source_quote=kwargs.get("source_quote"),
-                    confidence=kwargs.get("confidence", 1.0),
-                    metadata=metadata or {},
-                    first_seen=datetime.utcnow().isoformat(),
-                    last_updated=datetime.utcnow().isoformat(),
-                    parent_entity_id=kwargs.get("parent_entity_id"),
-                    used_entities=list(kwargs.get("used_entities", [])),
-                )
-                entry.checksum = compute_checksum(entry)
+            self.logger.error(
+                "Failed to track entity '%s' (transaction rolled back): %s. "
+                "Returning pre-failure state (%s).",
+                entity_id,
+                e,
+                "existing entry" if existing else "None (no prior entry existed)",
+                exc_info=True,
+            )
+            if existing is not None:
+                return copy.deepcopy(existing)
+            return None
 
         return entry
     
@@ -482,7 +481,8 @@ class ProvenanceManager:
                         entity_metadata = {**metadata, **entity.get("metadata", {})}
                         
                         try:
-                            self.track_entity(entity_id, source, entity_metadata, _conn=conn)
+                            with self.storage.savepoint(conn):
+                                self.track_entity(entity_id, source, entity_metadata, _conn=conn)
                             batch_count += 1
                         except Exception:
                             pass  # Continue with other entities in this batch
@@ -527,16 +527,17 @@ class ProvenanceManager:
                             continue
                         
                         try:
-                            self.track_chunk(
-                                chunk_id=chunk_id,
-                                source_document=source_document,
-                                source_path=source_path,
-                                start_index=chunk.get("start_index", 0),
-                                end_index=chunk.get("end_index", 0),
-                                parent_chunk_id=chunk.get("parent_chunk_id"),
-                                _conn=conn,
-                                **{**metadata, **chunk.get("metadata", {})}
-                            )
+                            with self.storage.savepoint(conn):
+                                self.track_chunk(
+                                    chunk_id=chunk_id,
+                                    source_document=source_document,
+                                    source_path=source_path,
+                                    start_index=chunk.get("start_index", 0),
+                                    end_index=chunk.get("end_index", 0),
+                                    parent_chunk_id=chunk.get("parent_chunk_id"),
+                                    _conn=conn,
+                                    **{**metadata, **chunk.get("metadata", {})}
+                                )
                             batch_count += 1
                         except Exception:
                             pass

semantica/provenance/storage.py

@@ -23,6 +23,8 @@ from abc import ABC, abstractmethod
 from typing import List, Optional, Dict, Any
 import sqlite3
 import json
+import uuid
+import threading
 from collections import deque
 from contextlib import contextmanager
 
@@ -108,6 +110,11 @@ class ProvenanceStorage(ABC):
         """Default transaction context manager for storage backends."""
         yield None
 
+    @contextmanager
+    def savepoint(self, conn: Any = None):
+        """Default savepoint context manager for nested transactions / per-item isolation."""
+        yield conn
+
     def _store_with_conn(self, conn: Any, entry: ProvenanceEntry) -> None:
         """Internal store method using an active connection/transaction."""
         self.store(entry)
@@ -142,6 +149,12 @@ class InMemoryStorage(ProvenanceStorage):
     def __init__(self):
         """Initialize in-memory storage."""
         self._entries: Dict[str, ProvenanceEntry] = {}
+        self._local = threading.local()
+
+    def _get_pending_stack(self) -> list:
+        if not hasattr(self._local, "pending_stack"):
+            self._local.pending_stack = []
+        return self._local.pending_stack
     
     def store(self, entry: ProvenanceEntry) -> None:
         """
@@ -235,15 +248,42 @@ class InMemoryStorage(ProvenanceStorage):
 
     @contextmanager
     def transaction(self):
-        """In-memory transaction context manager."""
-        yield None
+        """In-memory transaction context manager with staging buffer rollback."""
+        stack = self._get_pending_stack()
+        stack.append({})
+        try:
+            yield "IN_MEMORY_TX"
+            pending = stack.pop()
+            if stack:
+                stack[-1].update(pending)
+            else:
+                for entry in pending.values():
+                    self.store(entry)
+        except Exception:
+            if stack:
+                stack.pop()
+            raise
+
+    @contextmanager
+    def savepoint(self, conn: Any = None):
+        """In-memory savepoint context manager with staging buffer rollback."""
+        with self.transaction() as tx:
+            yield tx
 
     def _store_with_conn(self, conn: Any, entry: ProvenanceEntry) -> None:
         """Internal store method using an active connection/transaction."""
-        self.store(entry)
+        stack = self._get_pending_stack()
+        if stack:
+            stack[-1][entry.entity_id] = entry
+        else:
+            self.store(entry)
 
     def _retrieve_with_conn(self, conn: Any, entity_id: str) -> Optional[ProvenanceEntry]:
         """Internal retrieve method using an active connection/transaction."""
+        stack = self._get_pending_stack()
+        for pending in reversed(stack):
+            if entity_id in pending:
+                return pending[entity_id]
         return self.retrieve(entity_id)
 
 
@@ -348,6 +388,32 @@ class SQLiteStorage(ProvenanceStorage):
         finally:
             conn.close()
 
+    @contextmanager
+    def savepoint(self, conn: Any = None):
+        """
+        Context manager providing a SAVEPOINT for per-item rollback isolation
+        within an existing transaction.
+        
+        If no connection is provided, falls back to a full transaction.
+        """
+        if conn is None:
+            with self.transaction() as tx:
+                yield tx
+            return
+
+        sp_name = f"sp_{uuid.uuid4().hex}"
+        conn.execute(f"SAVEPOINT {sp_name}")
+        try:
+            yield conn
+            conn.execute(f"RELEASE {sp_name}")
+        except Exception:
+            try:
+                conn.execute(f"ROLLBACK TO {sp_name}")
+                conn.execute(f"RELEASE {sp_name}")
+            except Exception:
+                pass
+            raise
+
     @contextmanager
     def _read_connection(self):
         """

tests/provenance/test_manager.py

@@ -8,6 +8,7 @@ chunk tracking, source tracking, and lineage tracing.
 import pytest
 from unittest.mock import patch
 from semantica.provenance import ProvenanceManager, SourceReference
+from semantica.provenance.storage import InMemoryStorage, SQLiteStorage
 
 
 class TestProvenanceManager:
@@ -600,13 +601,14 @@ class TestProvenanceManager:
         assert "e_broken_child -> non_existent_parent" in check_broken["missing_references"]
 
     def test_track_entity_storage_error_swallowed(self):
-        """Test that track_entity swallows storage.store() exceptions and still returns a ProvenanceEntry."""
+        """Test that track_entity returns None when storage.store() fails on a brand-new entity,
+        since nothing was persisted (#782)."""
         prov_mgr = ProvenanceManager()
         with patch.object(prov_mgr.storage, "store", side_effect=RuntimeError("storage error")):
             entry = prov_mgr.track_entity("e_test", source="doc_1")
-            assert entry is not None
-            assert entry.entity_id == "e_test"
-            assert entry.checksum is not None
+            # Returning None is correct because nothing was actually persisted,
+            # and the old assertion was encoding the bug this issue was filed to fix.
+            assert entry is None
 
     def test_track_relationship_storage_error_swallowed(self):
         """Test that track_relationship swallows storage.store() exceptions and still returns a ProvenanceEntry."""
@@ -649,17 +651,226 @@ class TestProvenanceManager:
 
     def test_track_entity_pre_build_failure_fallback_skips_store(self):
         """Test that when track_entity fails before the entry is built (e.g. a
-        retrieve error inside the atomic transaction), the fallback entry only
-        gets a checksum and is not persisted via a second, out-of-transaction
-        storage.store() call (#784)."""
+        retrieve error inside the atomic transaction) on a brand-new entity,
+        it returns None (#782/#784)."""
         prov_mgr = ProvenanceManager()
         with patch.object(
             prov_mgr.storage, "_retrieve_with_conn", side_effect=RuntimeError("retrieve error")
         ), patch.object(prov_mgr.storage, "store") as mock_store:
             entry = prov_mgr.track_entity("e_test", source="doc_1")
-            assert entry is not None
-            assert entry.entity_id == "e_test"
-            assert entry.checksum is not None
+            # Returning None is correct because nothing was actually persisted,
+            # and the old assertion was encoding the bug this issue was filed to fix.
+            assert entry is None
             mock_store.assert_not_called()
 
+    @pytest.mark.parametrize("backend_type", ["memory", "sqlite"])
+    def test_track_entity_atomic_rollback_on_history_write_failure(self, backend_type, tmp_path):
+        """Test that history write failure rolls back transaction cleanly on both InMemory and SQLite storage."""
+        if backend_type == "memory":
+            class FlakyInMemory(InMemoryStorage):
+                def __init__(self):
+                    super().__init__()
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakyInMemory()
+        else:
+            db_path = str(tmp_path / "test_hist.db")
+            class FlakySQLite(SQLiteStorage):
+                def __init__(self, path):
+                    super().__init__(path)
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakySQLite(db_path)
+
+        mgr = ProvenanceManager(storage=storage)
+        v1 = mgr.track_entity("e1", source="doc1")
+        assert v1.source_document == "doc1"
+
+        storage.fail_on_call = 2
+        # Standalone call should not raise, should return pre-failure state (v1)
+        res = mgr.track_entity("e1", source="doc2")
+        in_storage = storage.retrieve("e1")
+        assert res.source_document == "doc1"
+        assert in_storage.source_document == "doc1"
+        assert [e.entity_id for e in storage.retrieve_all()] == ["e1"]
+
+        # Batch call (_conn supplied) should raise
+        storage.store_calls = 1
+        with pytest.raises(RuntimeError):
+            with mgr._get_or_create_transaction() as conn:
+                mgr.track_entity("e1", source="doc2", _conn=conn)
+
+    @pytest.mark.parametrize("backend_type", ["memory", "sqlite"])
+    def test_track_entity_atomic_rollback_on_primary_write_failure(self, backend_type, tmp_path):
+        """Test that primary write failure after history write rolls back the entire transaction."""
+        if backend_type == "memory":
+            class FlakyInMemory(InMemoryStorage):
+                def __init__(self):
+                    super().__init__()
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakyInMemory()
+        else:
+            db_path = str(tmp_path / "test_prim.db")
+            class FlakySQLite(SQLiteStorage):
+                def __init__(self, path):
+                    super().__init__(path)
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakySQLite(db_path)
+
+        mgr = ProvenanceManager(storage=storage)
+        v1 = mgr.track_entity("e1", source="doc1")
+        assert v1.source_document == "doc1"
+
+        storage.fail_on_call = 3
+        # Standalone call should not raise, should return pre-failure state (v1)
+        res = mgr.track_entity("e1", source="doc2")
+        in_storage = storage.retrieve("e1")
+        assert res.source_document == "doc1"
+        assert in_storage.source_document == "doc1"
+        assert [e.entity_id for e in storage.retrieve_all()] == ["e1"]
+
+        # Batch call (_conn supplied) should raise
+        storage.store_calls = 1
+        with pytest.raises(RuntimeError):
+            with mgr._get_or_create_transaction() as conn:
+                mgr.track_entity("e1", source="doc2", _conn=conn)
+
+    @pytest.mark.parametrize("backend_type", ["memory", "sqlite"])
+    def test_track_entity_rollback_returns_safe_copy(self, backend_type, tmp_path):
+        """Test that after a rollback returning the existing entry, mutating the returned object does not affect storage."""
+        if backend_type == "memory":
+            class FlakyInMemory(InMemoryStorage):
+                def __init__(self):
+                    super().__init__()
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakyInMemory()
+        else:
+            db_path = str(tmp_path / "test_safe_copy.db")
+            class FlakySQLite(SQLiteStorage):
+                def __init__(self, path):
+                    super().__init__(path)
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakySQLite(db_path)
+
+        mgr = ProvenanceManager(storage=storage)
+        v1 = mgr.track_entity("e1", source="doc1", metadata={"key": "val"})
+        assert v1.source_document == "doc1"
+
+        # Force failure on update
+        storage.fail_on_call = 2
+        res = mgr.track_entity("e1", source="doc2")
+        assert res.source_document == "doc1"
+
+        # Mutate the returned pre-failure object
+        res.source_document = "mutated"
+        res.metadata["tampered"] = True
+
+        # Assert stored copy in storage is completely unchanged
+        in_storage = storage.retrieve("e1")
+        assert in_storage.source_document == "doc1"
+        assert in_storage.metadata == {"key": "val"}
+        assert "tampered" not in in_storage.metadata
+
+    @pytest.mark.parametrize("backend_type", ["memory", "sqlite"])
+    def test_track_entities_batch_per_item_savepoint_rollback(self, backend_type, tmp_path):
+        """Test that in batch mode, an exception during primary write rolls back any staged history archive for that item."""
+        if backend_type == "memory":
+            class FlakyInMemory(InMemoryStorage):
+                def __init__(self):
+                    super().__init__()
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakyInMemory()
+        else:
+            db_path = str(tmp_path / "test_batch_sp.db")
+            class FlakySQLite(SQLiteStorage):
+                def __init__(self, path):
+                    super().__init__(path)
+                    self.store_calls = 0
+                    self.fail_on_call = None
+                def _store_with_conn(self, conn, entry):
+                    self.store_calls += 1
+                    if self.store_calls == self.fail_on_call:
+                        raise RuntimeError(f"Simulated error on call {self.store_calls}")
+                    super()._store_with_conn(conn, entry)
+            storage = FlakySQLite(db_path)
+
+        mgr = ProvenanceManager(storage=storage)
+        v1 = mgr.track_entity("e1", source="doc1")
+        assert v1.source_document == "doc1"
+
+        # Fail on call 3 (the primary update for e1, after its history entry was stored on call 2)
+        storage.fail_on_call = 3
+        count = mgr.track_entities_batch(
+            [{"id": "e1"}, {"id": "e2"}],
+            source="doc_batch"
+        )
+
+        assert count == 1  # Only e2 should succeed
+        all_entries = {e.entity_id: e for e in storage.retrieve_all()}
+
+        # Assert e2 is present
+        assert "e2" in all_entries
+        assert all_entries["e2"].source_document == "doc_batch"
+
+        # Assert e1 is untouched and NO history archive was left behind
+        assert all_entries["e1"].source_document == "doc1"
+        assert not any(":v:" in eid for eid in all_entries.keys()), f"Found leaked history entry: {list(all_entries.keys())}"
+
+    def test_track_entity_logs_exception_with_exc_info(self):
+        """Test that track_entity logs errors with exc_info=True when a storage write fails."""
+        class FlakyInMemory(InMemoryStorage):
+            def _store_with_conn(self, conn, entry):
+                raise RuntimeError("Simulated write failure")
+        storage = FlakyInMemory()
+        mgr = ProvenanceManager(storage=storage)
+
+        with patch.object(mgr.logger, "error") as mock_error:
+            res = mgr.track_entity("e1", source="doc1")
+            assert res is None
+            assert mock_error.called
+            _, kwargs = mock_error.call_args
+            assert kwargs.get("exc_info") is True
+
+
+
 

tests/provenance/test_sqlite_storage_performance_807.py

@@ -244,15 +244,17 @@ def test_chunks_batch_does_not_count_individually_failed_items(tmp_path):
 
 
 def test_track_entity_standalone_call_still_degrades_gracefully(tmp_path):
-    """Test that a direct (non-batch) track_entity() call still returns an
-    entry on storage failure instead of raising, preserving the public API's
-    existing graceful-degradation contract."""
+    """Test that a direct (non-batch) track_entity() call returns None
+    on storage failure for a brand-new entity instead of raising (#782),
+    preserving the public API's existing graceful-degradation contract."""
     db_path = str(tmp_path / "test_standalone_degrade.db")
     mgr = ProvenanceManager(storage_path=db_path)
 
     entry = mgr.track_entity("entity_1", source="doc_1", metadata={"bad": {1, 2, 3}})
 
-    assert entry.entity_id == "entity_1"
+    # Returning None is correct because nothing was actually persisted,
+    # and the old assertion was encoding the bug this issue was filed to fix.
+    assert entry is None
     assert mgr.storage.retrieve("entity_1") is None