Coverage for src/lexigram/admin/integrations/search_query.py: 96%

1from __future__ import annotations 

2 

3from dataclasses import replace 

4from typing import TYPE_CHECKING, Any 

5 

6from lexigram.admin.data.data_source import QueryResult 

7from lexigram.logging import get_logger 

8 

9if TYPE_CHECKING: 

10 from lexigram.admin.data.query import QuerySpec 

11 

12_log = get_logger(__name__) 

13 

14 

15class SearchQueryDataSourceWrapper: 

16 """Wraps an IDataSource to route within-resource search through FTS. 

17 

18 Intercepts ``find_many()`` and ``count()`` when the ``QuerySpec`` has 

19 an active search term. Instead of passing ``search``/``search_fields`` 

20 through to the inner data source (which would generate ILIKE), it 

21 queries the search engine for matching document IDs and adds an 

22 ``id__in`` filter. 

23 

24 Errors from the search engine are logged at DEBUG level and the query 

25 falls through without search filtering (no FTS, no ILIKE). 

26 """ 

27 

28 def __init__( 

29 self, 

30 inner: Any, 

31 search_engine: Any, 

32 index_name: str, 

33 fallback_to_like: bool = True, 

34 ) -> None: 

35 self._inner = inner 

36 self._search_engine = search_engine 

37 self._index_name = index_name 

38 self._fallback_to_like = fallback_to_like 

39 

40 async def find_one(self, item_id: Any) -> Any | None: 

41 return await self._inner.find_one(item_id) 

42 

43 async def find_many(self, query: QuerySpec) -> QueryResult: 

44 if not query.has_search: 

45 return await self._inner.find_many(query) 

46 

47 try: 

48 matched_ids, total = await self._fetch_matching_ids(query) 

49 _log.info( 

50 "search.query_wrapper_engine_result", 

51 index=self._index_name, 

52 search=query.search, 

53 matched_ids=len(matched_ids) if matched_ids else 0, 

54 total=total, 

55 fallback_to_like=self._fallback_to_like, 

56 ) 

57 if matched_ids is None: 

58 _log.info("search.query_wrapper_engine_error", index=self._index_name) 

59 cleared = replace(query, search=None, search_fields=[]) 

60 return await self._inner.find_many(cleared) 

61 if not matched_ids: 

62 if self._fallback_to_like: 

63 return await self._inner.find_many(query) 

64 return QueryResult(items=[], total=0) 

65 page_size = max(query.per_page, 1) 

66 page_offset = max(query.page, 1) - 1 

67 page_ids = matched_ids[ 

68 page_offset * page_size : (page_offset + 1) * page_size 

69 ] 

70 filtered = replace(query, search=None, search_fields=[]) 

71 if page_ids: 

72 filtered = filtered.with_filters(id__in=page_ids) 

73 return await self._inner.find_many(filtered) 

74 except Exception: 

75 _log.info( 

76 "search.query_wrapper_error", 

77 index=self._index_name, 

78 exc_info=True, 

79 ) 

80 cleared = replace(query, search=None, search_fields=[]) 

81 return await self._inner.find_many(cleared) 

82 

83 async def count(self, query: QuerySpec) -> int: 

84 if not query.has_search: 

85 return await self._inner.count(query) 

86 try: 

87 _, total = await self._fetch_matching_ids(query) 

88 return total 

89 except Exception: 

90 _log.debug( 

91 "search.count_wrapper_error", 

92 index=self._index_name, 

93 exc_info=True, 

94 ) 

95 cleared = replace(query, search=None, search_fields=[]) 

96 return await self._inner.count(cleared) 

97 

98 async def create(self, data: dict[str, Any]) -> Any: 

99 return await self._inner.create(data) 

100 

101 async def update(self, item_id: Any, data: dict[str, Any]) -> Any: 

102 return await self._inner.update(item_id, data) 

103 

104 async def delete(self, item_id: Any) -> bool: 

105 return await self._inner.delete(item_id) 

106 

107 async def bulk_create(self, items: list[dict[str, Any]]) -> list[Any]: 

108 return await self._inner.bulk_create(items) 

109 

110 async def bulk_update(self, ids: list[Any], data: dict[str, Any]) -> int: 

111 return await self._inner.bulk_update(ids, data) 

112 

113 async def bulk_delete(self, ids: list[Any]) -> int: 

114 return await self._inner.bulk_delete(ids) 

115 

116 async def _fetch_matching_ids( 

117 self, query: QuerySpec 

118 ) -> tuple[list[str] | None, int]: 

119 """Query the search engine and return (ids_for_page, total_count).""" 

120 page = max(query.page, 1) 

121 per_page = max(query.per_page, 1) 

122 offset = (page - 1) * per_page 

123 limit = per_page 

124 

125 result = await self._search_engine.search( 

126 index_name=self._index_name, 

127 query=query.search, 

128 limit=limit, 

129 offset=offset, 

130 ) 

131 

132 if hasattr(result, "is_ok"): 

133 if not result.is_ok(): 

134 return (None, 0) 

135 response = result.unwrap() 

136 else: 

137 response = result 

138 

139 total = response.total if hasattr(response, "total") else 0 

140 results_list = ( 

141 response.results if hasattr(response, "results") else (response or []) 

142 ) 

143 

144 ids = [] 

145 for r in results_list: 

146 rid = ( 

147 r.id 

148 if hasattr(r, "id") 

149 else (r.get("id") if isinstance(r, dict) else None) 

150 ) 

151 if rid is not None: 

152 ids.append(str(rid)) 

153 

154 return (ids, total) 

155 

156 

157def _empty_result(query: QuerySpec) -> QueryResult: 

158 from lexigram.admin.data.data_source import QueryResult 

159 

160 return QueryResult( 

161 items=[], 

162 total=0, 

163 page=query.page, 

164 per_page=query.per_page, 

165 ) 

166 

167 

168__all__ = ["SearchQueryDataSourceWrapper"]