Skip to content

Commit 9502e3d

Browse files
committed
fix(storage): close probe table handles in by-model search helpers
open_embeddings_table returns a live handle used only to resolve table_name before search_* reopens by name; release it in finally to avoid LanceDB handle leaks (PR #359 review).
1 parent 8bc2f18 commit 9502e3d

1 file changed

Lines changed: 62 additions & 34 deletions

File tree

src/xagent/core/tools/core/RAG_tools/storage/contracts.py

Lines changed: 62 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,22 @@
2929

3030
logger = logging.getLogger(__name__)
3131

32+
33+
def _release_embeddings_table_probe(table: Any) -> None:
34+
"""Close a table handle opened only to resolve embeddings table name.
35+
36+
``open_embeddings_table`` returns ``(table, name)``; by-model helpers then
37+
delegate to ``search_*`` methods that open their own handles by name. The
38+
probe handle must not be left open (LanceDB native tables support
39+
``close()``).
40+
"""
41+
if table is not None and hasattr(table, "close"):
42+
try:
43+
table.close()
44+
except Exception:
45+
pass
46+
47+
3248
# Field name whitelist for filter validation
3349
# Derived from all LanceDB table schemas in schema_manager.py
3450
_VALID_FILTER_FIELDS = frozenset(
@@ -877,15 +893,18 @@ def search_vectors_by_model(
877893
- metadata: Additional metadata
878894
"""
879895
_table, table_name = self.open_embeddings_table(model_tag)
880-
return self.search_vectors(
881-
table_name=table_name,
882-
query_vector=query_vector,
883-
top_k=top_k,
884-
filters=filters,
885-
vector_column_name=vector_column_name,
886-
user_id=user_id,
887-
is_admin=is_admin,
888-
)
896+
try:
897+
return self.search_vectors(
898+
table_name=table_name,
899+
query_vector=query_vector,
900+
top_k=top_k,
901+
filters=filters,
902+
vector_column_name=vector_column_name,
903+
user_id=user_id,
904+
is_admin=is_admin,
905+
)
906+
finally:
907+
_release_embeddings_table_probe(_table)
889908

890909
def search_fts_by_model(
891910
self,
@@ -916,15 +935,18 @@ def search_fts_by_model(
916935
List of search result dictionaries (see search_fts).
917936
"""
918937
_table, table_name = self.open_embeddings_table(model_tag)
919-
return self.search_fts(
920-
table_name=table_name,
921-
query_text=query_text,
922-
top_k=top_k,
923-
filters=filters,
924-
text_column_name=text_column_name,
925-
user_id=user_id,
926-
is_admin=is_admin,
927-
)
938+
try:
939+
return self.search_fts(
940+
table_name=table_name,
941+
query_text=query_text,
942+
top_k=top_k,
943+
filters=filters,
944+
text_column_name=text_column_name,
945+
user_id=user_id,
946+
is_admin=is_admin,
947+
)
948+
finally:
949+
_release_embeddings_table_probe(_table)
928950

929951
# --- Async variants (Phase 1A Option C: Hybrid approach) ---
930952

@@ -994,15 +1016,18 @@ async def search_vectors_by_model_async(
9941016
- metadata: Additional metadata
9951017
"""
9961018
_table, table_name = self.open_embeddings_table(model_tag)
997-
return await self.search_vectors_async(
998-
table_name=table_name,
999-
query_vector=query_vector,
1000-
top_k=top_k,
1001-
filters=filters,
1002-
vector_column_name=vector_column_name,
1003-
user_id=user_id,
1004-
is_admin=is_admin,
1005-
)
1019+
try:
1020+
return await self.search_vectors_async(
1021+
table_name=table_name,
1022+
query_vector=query_vector,
1023+
top_k=top_k,
1024+
filters=filters,
1025+
vector_column_name=vector_column_name,
1026+
user_id=user_id,
1027+
is_admin=is_admin,
1028+
)
1029+
finally:
1030+
_release_embeddings_table_probe(_table)
10061031

10071032
@abstractmethod
10081033
async def search_fts_async(
@@ -1068,13 +1093,16 @@ async def search_fts_by_model_async(
10681093
DatabaseOperationError: If FTS index is not configured or search fails.
10691094
"""
10701095
_table, table_name = self.open_embeddings_table(model_tag)
1071-
return await self.search_fts_async(
1072-
table_name=table_name,
1073-
query_text=query_text,
1074-
top_k=top_k,
1075-
filters=filters,
1076-
text_column_name=text_column_name,
1077-
)
1096+
try:
1097+
return await self.search_fts_async(
1098+
table_name=table_name,
1099+
query_text=query_text,
1100+
top_k=top_k,
1101+
filters=filters,
1102+
text_column_name=text_column_name,
1103+
)
1104+
finally:
1105+
_release_embeddings_table_probe(_table)
10781106

10791107
@abstractmethod
10801108
async def iter_batches_async(

0 commit comments

Comments
 (0)