2121
2222import asyncio
2323import sys
24+ import time
2425from pathlib import Path
26+ from typing import Callable , Literal
27+
28+ # Mirrors ``ProgressStatus`` in ``progress.py``; kept local (rather than imported)
29+ # so this module never pays the ``rich`` cost at import time — see
30+ # ``_make_optimize_event``.
31+ _OptimizeStatus = Literal ["running" , "done" , "failed" ]
2532
2633# Single source of truth for the three Lance table names created by the flow.
2734# Keep in sync with ``search_lancedb.TABLES`` (the values there mirror these).
3138 "yamlconfigindex_yaml_config" ,
3239)
3340
41+
42+ def _make_optimize_event (
43+ * ,
44+ status : _OptimizeStatus ,
45+ elapsed_s : float | None = None ,
46+ ):
47+ """Build a ``ProgressEvent(kind="optimize", …)`` lazily (progress is parent-side).
48+
49+ ``lance_optimize`` runs in-process in the parent (called by
50+ ``pipeline._maybe_run_serialized_optimize`` and
51+ ``server.run_refresh_pipeline``); it routes progress to the renderer via the
52+ in-process ``on_progress`` callback — NOT via stderr (which would corrupt
53+ the Live region). The import is local so the flow (which imports
54+ ``LANCE_TABLE_NAMES`` at definition time) never pays the ``rich`` cost.
55+ """
56+ from java_codebase_rag .progress import ProgressEvent
57+
58+ return ProgressEvent (
59+ kind = "optimize" ,
60+ phase = None ,
61+ pass_ = None ,
62+ done = None ,
63+ total = None ,
64+ status = status ,
65+ elapsed_s = elapsed_s ,
66+ )
67+
3468# Commit conflicts are transient; a handful of exponential-backoff retries is
3569# enough because, post-flow, there are no concurrent writers — only successive
3670# optimize/compaction passes within this single serialized call can still
@@ -60,7 +94,12 @@ async def _list_table_names(db: object) -> set[str]:
6094 return set (await db .table_names ())
6195
6296
63- async def optimize_lance_tables (index_dir : Path , * , quiet : bool = False ) -> dict [str , str ]:
97+ async def optimize_lance_tables (
98+ index_dir : Path ,
99+ * ,
100+ quiet : bool = False ,
101+ on_progress : Callable | None = None ,
102+ ) -> dict [str , str ]:
64103 """Optimize all known Lance tables under *index_dir*, serially, with retry.
65104
66105 Runs ``table.optimize()`` for each name in :data:`LANCE_TABLE_NAMES` that
@@ -73,6 +112,13 @@ async def optimize_lance_tables(index_dir: Path, *, quiet: bool = False) -> dict
73112 index_dir: directory holding the Lance tables (the flow's LanceDB URI).
74113 quiet: when True, suppress the per-table success/skip info lines on
75114 stderr (errors are always logged).
115+ on_progress: optional in-process progress callback (the parent's
116+ renderer ``on_progress``). When given, emits
117+ ``ProgressEvent(kind="optimize", status="running")`` on entry and a
118+ terminal ``status="done"``/``"failed"`` event on exit (covers BOTH
119+ call sites: ``pipeline._maybe_run_serialized_optimize`` and
120+ ``server.run_refresh_pipeline``). In-process only — NEVER prints to
121+ stderr (that would corrupt the Live region).
76122
77123 Returns:
78124 Mapping of table name → status. Values are ``"ok"``, ``"skipped"``
@@ -82,67 +128,95 @@ async def optimize_lance_tables(index_dir: Path, *, quiet: bool = False) -> dict
82128 # not pay the lancedb import cost at flow-definition time.
83129 import lancedb
84130
131+ if on_progress is not None :
132+ on_progress (_make_optimize_event (status = "running" ))
133+ t0 = time .perf_counter ()
85134 results : dict [str , str ] = {}
86- db = await lancedb . connect_async ( str ( index_dir ))
135+ failed = False
87136 try :
137+ db = await lancedb .connect_async (str (index_dir ))
88138 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 } ,
104- file = sys .stderr ,
105- )
106- continue
107139 try :
108- table = await db . open_table ( name )
140+ existing = await _list_table_names ( db )
109141 except Exception as exc :
110- results [name ] = f"error: open failed: { exc }
111142 print (
112- f"java-codebase-rag: optimize: { name } { exc } ,
143+ f"java-codebase-rag: optimize: failed to list tables in "
144+ f"{ index_dir } { exc } ,
113145 file = sys .stderr ,
114146 )
115- continue
116-
117- last_exc : BaseException | None = None
118- for attempt in range (_MAX_ATTEMPTS ):
147+ failed = True
148+ return {name : f"error: list failed: { exc } for name in LANCE_TABLE_NAMES }
149+
150+ for name in LANCE_TABLE_NAMES :
151+ if name not in existing :
152+ results [name ] = "skipped"
153+ if not quiet :
154+ print (
155+ f"java-codebase-rag: optimize: { name } ,
156+ file = sys .stderr ,
157+ )
158+ continue
119159 try :
120- await table .optimize ()
121- last_exc = None
122- break
160+ table = await db .open_table (name )
123161 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 :
162+ results [name ] = f"error: open failed: { exc }
163+ failed = True
135164 print (
136- f"java-codebase-rag: optimize: { name } ok " ,
165+ f"java-codebase-rag: optimize: { name } open failed: { exc } ,
137166 file = sys .stderr ,
138167 )
139- else :
140- results [name ] = f"error: { last_exc }
141- print (
142- f"java-codebase-rag: optimize: { name } { last_exc } ,
143- file = sys .stderr ,
144- )
168+ continue
169+
170+ last_exc : BaseException | None = None
171+ for attempt in range (_MAX_ATTEMPTS ):
172+ try :
173+ await table .optimize ()
174+ last_exc = None
175+ break
176+ except Exception as exc :
177+ last_exc = exc
178+ if _is_retryable (exc ) and attempt < _MAX_ATTEMPTS - 1 :
179+ await asyncio .sleep (_BASE_BACKOFF_S * (2 ** attempt ))
180+ continue
181+ # Non-retryable, or retries exhausted: stop the loop and
182+ # surface below — do not swallow silently.
183+ break
184+
185+ if last_exc is None :
186+ results [name ] = "ok"
187+ if not quiet :
188+ print (
189+ f"java-codebase-rag: optimize: { name } ,
190+ file = sys .stderr ,
191+ )
192+ else :
193+ results [name ] = f"error: { last_exc }
194+ failed = True
195+ print (
196+ f"java-codebase-rag: optimize: { name } { last_exc } ,
197+ file = sys .stderr ,
198+ )
199+ finally :
200+ # ``AsyncConnection.close`` is a *sync* method in lancedb 0.30.x.
201+ db .close ()
202+ return results
203+ except Exception :
204+ # An unexpected exception (e.g. ``connect_async`` raised, or a table-
205+ # independent failure) must still flip the terminal event to failed so
206+ # the renderer's task doesn't render a green check on a crash. Re-raise
207+ # after marking — the caller (``_maybe_run_serialized_optimize`` /
208+ # ``run_refresh_pipeline``) treats optimize failure as non-fatal and
209+ # logs it, but the renderer must reflect the truth.
210+ failed = True
211+ raise
145212 finally :
146- # ``AsyncConnection.close`` is a *sync* method in lancedb 0.30.x.
147- db .close ()
148- return results
213+ # Always emit a terminal optimize event so the renderer's task never
214+ # hangs at "running" — even on exception (the parent treats a failed
215+ # optimize as non-fatal: the index is still searchable un-compacted).
216+ if on_progress is not None :
217+ on_progress (
218+ _make_optimize_event (
219+ status = "failed" if failed else "done" ,
220+ elapsed_s = time .perf_counter () - t0 ,
221+ )
222+ )
0 commit comments