|
| 1 | +"""Serialized post-flow LanceDB optimize with commit-conflict retry. |
| 2 | +
|
| 3 | +cocoindex 1.0.7 schedules ``table.optimize()`` (a LanceDB **Rewrite**/compaction |
| 4 | +transaction) as a *background* ``asyncio`` task that races concurrent |
| 5 | +``table.delete()`` (**Delete**) transactions emitted by later mutation batches. |
| 6 | +LanceDB does not allow a Rewrite to commit concurrently with a Delete |
| 7 | +(upstream lancedb#1504 — "We do not support concurrent deletes right now"), |
| 8 | +which surfaces as a flood of:: |
| 9 | +
|
| 10 | + RuntimeError: lance error: Retryable commit conflict for version N: \ |
| 11 | +This Rewrite transaction was preempted by concurrent transaction Delete ... |
| 12 | +
|
| 13 | +To eliminate the race, the flow (``java_index_flow_lancedb.py``) disables the |
| 14 | +in-flight background optimize entirely by raising |
| 15 | +``num_transactions_before_optimize`` to a value that is effectively never |
| 16 | +reached. This module then performs a *single*, serialized optimize after the |
| 17 | +flow returns (exit 0 → no concurrent writers), retrying the rare residual |
| 18 | +commit conflict that two internal compaction passes can still produce. |
| 19 | +""" |
| 20 | +from __future__ import annotations |
| 21 | + |
| 22 | +import asyncio |
| 23 | +import sys |
| 24 | +from pathlib import Path |
| 25 | + |
| 26 | +# Single source of truth for the three Lance table names created by the flow. |
| 27 | +# Keep in sync with ``search_lancedb.TABLES`` (the values there mirror these). |
| 28 | +LANCE_TABLE_NAMES: tuple[str, ...] = ( |
| 29 | + "javacodeindex_java_code", |
| 30 | + "sqlschemaindex_sql_schema", |
| 31 | + "yamlconfigindex_yaml_config", |
| 32 | +) |
| 33 | + |
| 34 | +# Commit conflicts are transient; a handful of exponential-backoff retries is |
| 35 | +# enough because, post-flow, there are no concurrent writers — only successive |
| 36 | +# optimize/compaction passes within this single serialized call can still |
| 37 | +# transiently preempt one another. |
| 38 | +_MAX_ATTEMPTS = 6 |
| 39 | +_BASE_BACKOFF_S = 0.1 |
| 40 | + |
| 41 | +# Substrings identifying the retryable Lance commit-conflict error. LanceDB |
| 42 | +# wraps the underlying lance error text into the raised ``RuntimeError`` str, |
| 43 | +# so a substring match is the robust detector (no dedicated exception type). |
| 44 | +_RETRYABLE_MARKERS = ( |
| 45 | + "Retryable commit conflict", |
| 46 | + "preempted by concurrent transaction", |
| 47 | +) |
| 48 | + |
| 49 | + |
| 50 | +def _is_retryable(exc: BaseException) -> bool: |
| 51 | + text = str(exc) |
| 52 | + return any(marker in text for marker in _RETRYABLE_MARKERS) |
| 53 | + |
| 54 | + |
| 55 | +async def _list_table_names(db: object) -> set[str]: |
| 56 | + """Existing table names across LanceDB API variants (``list_tables`` ≥ ``table_names``).""" |
| 57 | + if hasattr(db, "list_tables"): |
| 58 | + response = await db.list_tables() |
| 59 | + return set(getattr(response, "tables", response)) |
| 60 | + return set(await db.table_names()) |
| 61 | + |
| 62 | + |
| 63 | +async def optimize_lance_tables(index_dir: Path, *, quiet: bool = False) -> dict[str, str]: |
| 64 | + """Optimize all known Lance tables under *index_dir*, serially, with retry. |
| 65 | +
|
| 66 | + Runs ``table.optimize()`` for each name in :data:`LANCE_TABLE_NAMES` that |
| 67 | + exists in the DB. Retryable commit conflicts are retried with exponential |
| 68 | + backoff; any other exception (or an exhausted retry budget) is captured |
| 69 | + per-table in the returned dict and logged to **stderr** — never stdout, |
| 70 | + since this is callable from stdio-MCP / JSON-stdout contexts. |
| 71 | +
|
| 72 | + Args: |
| 73 | + index_dir: directory holding the Lance tables (the flow's LanceDB URI). |
| 74 | + quiet: when True, suppress the per-table success/skip info lines on |
| 75 | + stderr (errors are always logged). |
| 76 | +
|
| 77 | + Returns: |
| 78 | + Mapping of table name → status. Values are ``"ok"``, ``"skipped"`` |
| 79 | + (table absent — e.g. a repo with no SQL/YAML), or ``"error: <text>"``. |
| 80 | + """ |
| 81 | + # Lazy import: the flow imports this module for LANCE_TABLE_NAMES and must |
| 82 | + # not pay the lancedb import cost at flow-definition time. |
| 83 | + import lancedb |
| 84 | + |
| 85 | + results: dict[str, str] = {} |
| 86 | + db = await lancedb.connect_async(str(index_dir)) |
| 87 | + try: |
| 88 | + try: |
| 89 | + existing = await _list_table_names(db) |
| 90 | + except Exception as exc: |
| 91 | + print( |
| 92 | + f"java-codebase-rag: optimize: failed to list tables in " |
| 93 | + f"{index_dir}: {exc}", |
| 94 | + file=sys.stderr, |
| 95 | + ) |
| 96 | + return {name: f"error: list failed: {exc}" for name in LANCE_TABLE_NAMES} |
| 97 | + |
| 98 | + for name in LANCE_TABLE_NAMES: |
| 99 | + if name not in existing: |
| 100 | + results[name] = "skipped" |
| 101 | + if not quiet: |
| 102 | + print( |
| 103 | + f"java-codebase-rag: optimize: {name} absent, skipped", |
| 104 | + file=sys.stderr, |
| 105 | + ) |
| 106 | + continue |
| 107 | + try: |
| 108 | + table = await db.open_table(name) |
| 109 | + except Exception as exc: |
| 110 | + results[name] = f"error: open failed: {exc}" |
| 111 | + print( |
| 112 | + f"java-codebase-rag: optimize: {name} open failed: {exc}", |
| 113 | + file=sys.stderr, |
| 114 | + ) |
| 115 | + continue |
| 116 | + |
| 117 | + last_exc: BaseException | None = None |
| 118 | + for attempt in range(_MAX_ATTEMPTS): |
| 119 | + try: |
| 120 | + await table.optimize() |
| 121 | + last_exc = None |
| 122 | + break |
| 123 | + except Exception as exc: |
| 124 | + last_exc = exc |
| 125 | + if _is_retryable(exc) and attempt < _MAX_ATTEMPTS - 1: |
| 126 | + await asyncio.sleep(_BASE_BACKOFF_S * (2**attempt)) |
| 127 | + continue |
| 128 | + # Non-retryable, or retries exhausted: stop the loop and |
| 129 | + # surface below — do not swallow silently. |
| 130 | + break |
| 131 | + |
| 132 | + if last_exc is None: |
| 133 | + results[name] = "ok" |
| 134 | + if not quiet: |
| 135 | + print( |
| 136 | + f"java-codebase-rag: optimize: {name} ok", |
| 137 | + file=sys.stderr, |
| 138 | + ) |
| 139 | + else: |
| 140 | + results[name] = f"error: {last_exc}" |
| 141 | + print( |
| 142 | + f"java-codebase-rag: optimize: {name} failed: {last_exc}", |
| 143 | + file=sys.stderr, |
| 144 | + ) |
| 145 | + finally: |
| 146 | + # ``AsyncConnection.close`` is a *sync* method in lancedb 0.30.x. |
| 147 | + db.close() |
| 148 | + return results |
0 commit comments