superlocalmemory 3.8.2 → 3.8.5

package/src/superlocalmemory/core/graph_pruner.py

@@ -24,6 +24,10 @@ import time
 import uuid
 from datetime import datetime, UTC
 from pathlib import Path
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from superlocalmemory.storage.database import DatabaseManager
 
 logger = logging.getLogger("superlocalmemory.graph_pruner")
 
@@ -34,70 +38,118 @@ _CHAIN_BATCH_LIMIT = 10_000
 _MAX_DEGREE_PER_NODE: int = 100
 _HUB_PRUNE_BATCH: int = 500  # delete edges in batches to avoid giant IN clauses
 
+# Fix A: maximum rows to DELETE per transaction (keeps lock-hold < 10 ms).
+_BATCH_SIZE: int = 1_000
+# Brief sleep between batch-delete transactions so other writers get a turn.
+_BATCH_YIELD_S: float = 0.005
+
 
 def prune_graph(
-    db_path: str | Path,
+    db_or_path: "Union[DatabaseManager, str, Path]",
     profile_id: str = "default",
     dry_run: bool = False,
     cap_degree: bool = True,
+    max_degree: int = _MAX_DEGREE_PER_NODE,
+    min_edge_weight: float = 0.0,
 ) -> dict:
     """Run all graph pruning strategies for a specific profile.
 
+    Fix A: accepts a DatabaseManager as the first argument (preferred) to
+    route all writes through the process-level serialisation lock, eliminating
+    the lock-storm contribution from the legacy open-own-connection pattern.
+
+    For backward compatibility a ``str | Path`` is also accepted: a temporary
+    DatabaseManager is constructed from it with a deprecation warning.
+
+    Each DELETE strategy now runs in its own short transaction (≤ _BATCH_SIZE
+    rows) with a brief yield between batches.  This caps the WAL write-lock
+    hold to < 10 ms per batch instead of potentially 30+ seconds for the
+    full prune.
+
+    v3.8.4-G (#84): two new parameters for config-driven thinning:
+      - ``max_degree``: per-node in/out degree cap (default = _MAX_DEGREE_PER_NODE=100).
+      - ``min_edge_weight``: edges with weight strictly below this floor are
+        discarded before any other step.  Default 0.0 = no floor (current behaviour).
+
+    NOTE (Flaw 2 from CRIT): orphan IDs are read once upfront, then deleted
+    in batches.  New orphan edges created by the materialiser DURING a prune
+    run are NOT in the initial ID list and will be pruned on the NEXT cycle.
+    The ``total_after`` stat therefore reflects reality after the last batch
+    completes, which may be slightly different from total_before - removed.
+    This is acceptable — no data loss, just a stat nuance.
+
     Returns stats dict with counts for each strategy.
     """
-    conn = sqlite3.connect(str(db_path))
-    conn.execute("PRAGMA journal_mode=WAL")
-    conn.execute("PRAGMA busy_timeout=30000")
-    conn.row_factory = sqlite3.Row
+    from superlocalmemory.storage.database import DatabaseManager
 
-    stats = {
+    if isinstance(db_or_path, DatabaseManager):
+        db = db_or_path
+    else:
+        logger.warning(
+            "prune_graph: passing a db_path is deprecated — pass a "
+            "DatabaseManager instead (Fix A backward-compat shim active)"
+        )
+        db = DatabaseManager(db_or_path)
+
+    stats: dict = {
         "orphans_removed": 0,
         "supersedes_collapsed": 0,
         "self_loops_removed": 0,
         "duplicates_removed": 0,
         "hub_edges_removed": 0,
-        "association_orphans_removed": 0,  # gi-04
+        "association_orphans_removed": 0,
+        "low_weight_removed": 0,
         "total_before": 0,
         "total_after": 0,
     }
 
     try:
-        c = conn.cursor()
-        c.execute(
-            "SELECT COUNT(*) as cnt FROM graph_edges WHERE profile_id = ?",
+        before = db.execute(
+            "SELECT COUNT(*) AS cnt FROM graph_edges WHERE profile_id = ?",
             (profile_id,),
         )
-        stats["total_before"] = c.fetchone()["cnt"]
+        stats["total_before"] = int(before[0]["cnt"]) if before else 0
 
         start = time.time()
 
-        # Explicit transaction for atomicity
-        c.execute("BEGIN")
-
-        stats["orphans_removed"] = _remove_orphan_edges(c, profile_id, dry_run)
-        stats["self_loops_removed"] = _remove_self_loops(c, profile_id, dry_run)
-        stats["duplicates_removed"] = _remove_duplicate_edges(c, profile_id, dry_run)
-        stats["supersedes_collapsed"] = _collapse_supersedes_chains(
-            c, profile_id, dry_run,
-        )
-        if cap_degree:
-            stats["hub_edges_removed"] = _cap_node_degree(
-                c, profile_id, _MAX_DEGREE_PER_NODE, dry_run,
+        if not dry_run:
+            # v3.8.4-G: min_edge_weight floor — run FIRST so subsequent steps
+            # operate on the already-thinned edge set.
+            if min_edge_weight > 0.0:
+                stats["low_weight_removed"] = _remove_low_weight_edges_batched(
+                    db, profile_id, min_edge_weight,
+                )
+            stats["orphans_removed"] = _remove_orphan_edges_batched(
+                db, profile_id,
+            )
+            stats["self_loops_removed"] = _remove_self_loops_batched(
+                db, profile_id,
+            )
+            stats["duplicates_removed"] = _remove_duplicate_edges_batched(
+                db, profile_id,
+            )
+            stats["supersedes_collapsed"] = _collapse_supersedes_chains_batched(
+                db, profile_id,
+            )
+            if cap_degree:
+                stats["hub_edges_removed"] = _cap_node_degree_batched(
+                    db, profile_id, max_degree,
+                )
+            stats["association_orphans_removed"] = (
+                _remove_orphan_association_edges_batched(db, profile_id)
             )
-        stats["association_orphans_removed"] = _remove_orphan_association_edges(
-            c, profile_id, dry_run,
-        )
-
-        if dry_run:
-            c.execute("ROLLBACK")
         else:
-            c.execute("COMMIT")
+            # Dry-run: read-only estimates via a temporary connection
+            # (DatabaseManager is not used here — no writes)
+            stats.update(
+                _dry_run_counts(db.db_path, profile_id, cap_degree)
+            )
 
-        c.execute(
-            "SELECT COUNT(*) as cnt FROM graph_edges WHERE profile_id = ?",
+        after = db.execute(
+            "SELECT COUNT(*) AS cnt FROM graph_edges WHERE profile_id = ?",
             (profile_id,),
         )
-        stats["total_after"] = c.fetchone()["cnt"]
+        stats["total_after"] = int(after[0]["cnt"]) if after else 0
 
         elapsed = time.time() - start
         total_removed = stats["total_before"] - stats["total_after"]
@@ -116,16 +168,361 @@ def prune_graph(
     except Exception as exc:
         logger.error("Graph pruning failed: %s", exc, exc_info=True)
         stats["error"] = str(exc)
-        try:
-            c.execute("ROLLBACK")
-        except Exception:
-            pass
-    finally:
-        conn.close()
 
     return stats
 
 
+# ---------------------------------------------------------------------------
+# Batched delete helpers (Fix A — short transactions, bounded lock hold)
+# ---------------------------------------------------------------------------
+
+def _batch_delete_by_ids(
+    db: "DatabaseManager",
+    table: str,
+    ids: list,
+    id_col: str = "edge_id",
+) -> int:
+    """Delete rows in ``table`` where ``id_col`` IN ``ids``, _BATCH_SIZE at a time.
+
+    Each batch runs in its own transaction so the write lock is released
+    between batches. Returns total rows deleted.
+    """
+    removed = 0
+    for start in range(0, len(ids), _BATCH_SIZE):
+        batch = ids[start:start + _BATCH_SIZE]
+        ph = ",".join("?" * len(batch))
+        with db.transaction():
+            db.execute(
+                f"DELETE FROM {table} WHERE {id_col} IN ({ph})",
+                tuple(batch),
+            )
+        removed += len(batch)
+        if start + _BATCH_SIZE < len(ids):
+            time.sleep(_BATCH_YIELD_S)
+    return removed
+
+
+def _remove_orphan_edges_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+) -> int:
+    """Remove graph_edges whose source OR target no longer exists in facts/entities.
+
+    Read all orphan IDs upfront, then batch-delete.  New orphans created
+    during the run will be handled in the next maintenance cycle.
+    """
+    rows = db.execute(
+        """
+        SELECT edge_id FROM graph_edges
+        WHERE profile_id = ?
+          AND (
+            (source_id NOT IN (SELECT fact_id FROM atomic_facts)
+             AND source_id NOT IN (SELECT entity_id FROM canonical_entities))
+            OR
+            (target_id NOT IN (SELECT fact_id FROM atomic_facts)
+             AND target_id NOT IN (SELECT entity_id FROM canonical_entities))
+          )
+        """,
+        (profile_id,),
+    )
+    ids = [dict(r)["edge_id"] for r in rows]
+    if not ids:
+        return 0
+    return _batch_delete_by_ids(db, "graph_edges", ids)
+
+
+def _remove_self_loops_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+) -> int:
+    """Remove self-loop edges (source_id == target_id)."""
+    rows = db.execute(
+        "SELECT edge_id FROM graph_edges "
+        "WHERE source_id = target_id AND profile_id = ?",
+        (profile_id,),
+    )
+    ids = [dict(r)["edge_id"] for r in rows]
+    if not ids:
+        return 0
+    return _batch_delete_by_ids(db, "graph_edges", ids)
+
+
+def _remove_duplicate_edges_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+) -> int:
+    """Remove duplicate edges (same source+target+type), keeping highest weight."""
+    rows = db.execute(
+        """
+        SELECT edge_id FROM graph_edges ge_outer
+        WHERE profile_id = ?
+          AND edge_id NOT IN (
+            SELECT edge_id FROM graph_edges ge1
+            WHERE profile_id = ?
+              AND weight = (
+                SELECT MAX(weight) FROM graph_edges ge2
+                WHERE ge2.source_id = ge1.source_id
+                  AND ge2.target_id = ge1.target_id
+                  AND ge2.edge_type = ge1.edge_type
+                  AND ge2.profile_id = ge1.profile_id
+              )
+            GROUP BY source_id, target_id, edge_type
+          )
+        """,
+        (profile_id, profile_id),
+    )
+    ids = [dict(r)["edge_id"] for r in rows]
+    if not ids:
+        return 0
+    return _batch_delete_by_ids(db, "graph_edges", ids)
+
+
+def _collapse_supersedes_chains_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+) -> int:
+    """Collapse supersedes chains: A→B, B→C becomes A→C (B→C removed)."""
+    rows = db.execute(
+        """
+        SELECT e1.edge_id AS e1_id, e1.source_id AS a, e1.target_id AS b,
+               e1.weight AS e1_weight,
+               e2.edge_id AS e2_id, e2.target_id AS c
+        FROM graph_edges e1
+        JOIN graph_edges e2 ON e1.target_id = e2.source_id
+        WHERE e1.edge_type = 'supersedes'
+          AND e2.edge_type = 'supersedes'
+          AND e1.profile_id = ?
+          AND e2.profile_id = ?
+        LIMIT ?
+        """,
+        (profile_id, profile_id, _CHAIN_BATCH_LIMIT),
+    )
+    chains = [dict(r) for r in rows]
+    if not chains:
+        return 0
+
+    if len(chains) >= _CHAIN_BATCH_LIMIT:
+        logger.warning(
+            "Supersedes chain collapse hit limit (%d). "
+            "More chains may exist — will process in next cycle.",
+            _CHAIN_BATCH_LIMIT,
+        )
+
+    now = datetime.now(UTC).isoformat()
+    delete_ids: list[str] = []
+    insert_rows: list[tuple] = []
+
+    for chain in chains:
+        delete_ids.append(chain["e2_id"])
+        new_edge_id = uuid.uuid4().hex[:16]
+        insert_rows.append((
+            new_edge_id, profile_id,
+            chain["a"], chain["c"],
+            "supersedes", chain["e1_weight"] or 1.0, now,
+        ))
+
+    # Delete B→C edges in batches, then insert A→C shortcuts
+    removed = _batch_delete_by_ids(db, "graph_edges", delete_ids)
+
+    for start in range(0, len(insert_rows), _BATCH_SIZE):
+        batch = insert_rows[start:start + _BATCH_SIZE]
+        with db.transaction():
+            for row in batch:
+                db.execute(
+                    "INSERT OR IGNORE INTO graph_edges "
+                    "(edge_id, profile_id, source_id, target_id, "
+                    " edge_type, weight, created_at) "
+                    "VALUES (?,?,?,?,?,?,?)",
+                    row,
+                )
+        if start + _BATCH_SIZE < len(insert_rows):
+            time.sleep(_BATCH_YIELD_S)
+
+    return removed
+
+
+def _cap_node_degree_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+    max_degree: int,
+) -> int:
+    """Remove low-weight edges from hub nodes (degree > max_degree).
+
+    Uses the same window-function approach as the legacy cursor version but
+    selects IDs into Python then batch-deletes through DatabaseManager.
+    """
+    rows = db.execute(
+        """
+        SELECT edge_id FROM (
+            SELECT edge_id,
+                   ROW_NUMBER() OVER (
+                       PARTITION BY source_id ORDER BY weight DESC
+                   ) AS out_rn,
+                   ROW_NUMBER() OVER (
+                       PARTITION BY target_id ORDER BY weight DESC
+                   ) AS in_rn
+            FROM graph_edges
+            WHERE profile_id = ?
+        ) WHERE out_rn > ? OR in_rn > ?
+        """,
+        (profile_id, max_degree, max_degree),
+    )
+    ids = [dict(r)["edge_id"] for r in rows]
+    if not ids:
+        return 0
+    removed = _batch_delete_by_ids(db, "graph_edges", ids)
+    logger.info(
+        "_cap_node_degree_batched: removed %d low-weight edges (max_degree=%d)",
+        removed, max_degree,
+    )
+    return removed
+
+
+def _remove_low_weight_edges_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+    min_weight: float,
+) -> int:
+    """Remove graph_edges with weight strictly below ``min_weight``.
+
+    v3.8.4-G (#84): implements the ``min_edge_weight`` floor from
+    ``GraphPruningConfig``.  Runs before orphan/duplicate/degree pruning so
+    subsequent steps operate on the already-thinned set.
+
+    A ``min_weight`` of 0.0 means no floor — no edges are touched.  The
+    caller in ``prune_graph()`` guards the call with ``if min_edge_weight > 0.0``
+    so this function is only invoked when there is an actual floor to enforce.
+    """
+    rows = db.execute(
+        "SELECT edge_id FROM graph_edges WHERE profile_id = ? AND weight < ?",
+        (profile_id, min_weight),
+    )
+    ids = [dict(r)["edge_id"] for r in rows]
+    if not ids:
+        return 0
+    removed = _batch_delete_by_ids(db, "graph_edges", ids)
+    logger.info(
+        "_remove_low_weight_edges_batched: removed %d edges (min_weight=%.4f)",
+        removed, min_weight,
+    )
+    return removed
+
+
+def _remove_orphan_association_edges_batched(
+    db: "DatabaseManager",
+    profile_id: str,
+) -> int:
+    """Remove association_edges whose source/target fact no longer exists.
+
+    Fix: PK column is ``edge_id``, not ``id`` (schema_v32.py line 177).
+    """
+    rows = db.execute(
+        """
+        SELECT edge_id FROM association_edges
+        WHERE profile_id = ?
+          AND (
+            source_fact_id NOT IN (
+                SELECT fact_id FROM atomic_facts WHERE profile_id = ?
+            )
+            OR target_fact_id NOT IN (
+                SELECT fact_id FROM atomic_facts WHERE profile_id = ?
+            )
+          )
+        """,
+        (profile_id, profile_id, profile_id),
+    )
+    ids = [dict(r)["edge_id"] for r in rows]
+    if not ids:
+        return 0
+    return _batch_delete_by_ids(db, "association_edges", ids, id_col="edge_id")
+
+
+# ---------------------------------------------------------------------------
+# Dry-run estimates (read-only; uses a direct connection to avoid side effects)
+# ---------------------------------------------------------------------------
+
+def _dry_run_counts(
+    db_path: Path,
+    profile_id: str,
+    cap_degree: bool,
+) -> dict:
+    """Return estimated removal counts without modifying anything."""
+    counts: dict = {
+        "orphans_removed": 0,
+        "supersedes_collapsed": 0,
+        "self_loops_removed": 0,
+        "duplicates_removed": 0,
+        "hub_edges_removed": 0,
+        "association_orphans_removed": 0,
+    }
+    try:
+        conn = sqlite3.connect(str(db_path))
+        conn.execute("PRAGMA busy_timeout=10000")
+        conn.row_factory = sqlite3.Row
+        c = conn.cursor()
+
+        c.execute(
+            """
+            SELECT COUNT(*) AS cnt FROM graph_edges
+            WHERE profile_id = ?
+              AND (
+                (source_id NOT IN (SELECT fact_id FROM atomic_facts)
+                 AND source_id NOT IN (SELECT entity_id FROM canonical_entities))
+                OR
+                (target_id NOT IN (SELECT fact_id FROM atomic_facts)
+                 AND target_id NOT IN (SELECT entity_id FROM canonical_entities))
+              )
+            """,
+            (profile_id,),
+        )
+        counts["orphans_removed"] = c.fetchone()["cnt"]
+
+        c.execute(
+            "SELECT COUNT(*) AS cnt FROM graph_edges "
+            "WHERE source_id = target_id AND profile_id = ?",
+            (profile_id,),
+        )
+        counts["self_loops_removed"] = c.fetchone()["cnt"]
+
+        c.execute(
+            """
+            SELECT
+                (SELECT COUNT(*) FROM graph_edges WHERE profile_id = ?) -
+                (SELECT COUNT(*) FROM (
+                    SELECT source_id, target_id, edge_type
+                    FROM graph_edges WHERE profile_id = ?
+                    GROUP BY source_id, target_id, edge_type
+                )) AS cnt
+            """,
+            (profile_id, profile_id),
+        )
+        counts["duplicates_removed"] = max(c.fetchone()["cnt"] or 0, 0)
+
+        if cap_degree:
+            c.execute(
+                """
+                SELECT COUNT(*) AS cnt FROM (
+                    SELECT edge_id,
+                           ROW_NUMBER() OVER (
+                               PARTITION BY source_id ORDER BY weight DESC
+                           ) AS out_rn,
+                           ROW_NUMBER() OVER (
+                               PARTITION BY target_id ORDER BY weight DESC
+                           ) AS in_rn
+                    FROM graph_edges
+                    WHERE profile_id = ?
+                ) WHERE out_rn > ? OR in_rn > ?
+                """,
+                (profile_id, _MAX_DEGREE_PER_NODE, _MAX_DEGREE_PER_NODE),
+            )
+            counts["hub_edges_removed"] = c.fetchone()["cnt"]
+
+        conn.close()
+    except Exception as exc:
+        logger.warning("dry-run counts failed: %s", exc)
+    return counts
+
+
 def _remove_orphan_association_edges(
     c: sqlite3.Cursor,
     profile_id: str,