feat(search): add lib/search crate (client, index, query, sync, metrics)
zhenyi committed
Added lib/search/Cargo.toml lib/search/Cargo.toml +26 −0 View lib/search/Cargo.toml
@@ -0,0 +1,26 @@
1[package]
2name = "search"
3version.workspace = true
4edition.workspace = true
5authors.workspace = true
6categories.workspace = true
7keywords.workspace = true
8repository.workspace = true
9homepage.workspace = true
10
11[lib]
12path = "lib.rs"
13name = "search"
14
15[dependencies]
16# disable jwt features to avoid pulling aws-lc / jsonwebtoken; tenant tokens can re-enable later
17meilisearch-sdk = { version = "0.33", default-features = false, features = ["reqwest", "tls"] }
18config = { workspace = true }
19tokio = { workspace = true }
20serde = { workspace = true }
21serde_json = { workspace = true }
22thiserror = { workspace = true }
23tracing = { workspace = true }
24metrics = { workspace = true }
25lombok = { workspace = true }
26queue = { workspace = true }
No newline at end of file
Added lib/search/client.rs lib/search/client.rs +186 −0 View lib/search/client.rs
@@ -0,0 +1,186 @@
1use std::sync::Arc;
2use std::time::Duration;
3
4use meilisearch_sdk::client::{Client as MeiliClient, Health};
5use meilisearch_sdk::errors::ErrorCode;
6use meilisearch_sdk::task_info::TaskInfo;
7use meilisearch_sdk::tasks::Task;
8
9use crate::error::{SearchError, SearchResult};
10use crate::index::IndexHandle;
11use crate::metrics;
12use crate::sync::SyncedIndex;
13use config::search::AppSearchSettings;
14
15#[derive(Clone)]
16pub struct SearchClient {
17 inner: MeiliClient,
18 cfg: Arc<AppSearchSettings>,
19}
20
21impl SearchClient {
22 pub fn from_config(cfg: AppSearchSettings) -> SearchResult<Option<Self>> {
23 if !cfg.enabled {
24 return Ok(None);
25 }
26 cfg.is_valid().map_err(SearchError::Config)?;
27 let inner = MeiliClient::new(&cfg.host, cfg.api_key.as_deref())
28 .map_err(SearchError::Meilisearch)?;
29 Ok(Some(Self {
30 inner,
31 cfg: Arc::new(cfg),
32 }))
33 }
34
35 pub fn new(host: impl Into<String>, api_key: Option<impl Into<String>>) -> SearchResult<Self> {
36 let host = host.into();
37 let api_key = api_key.map(Into::into);
38 let inner =
39 MeiliClient::new(&host, api_key.as_deref()).map_err(SearchError::Meilisearch)?;
40 let mut cfg = AppSearchSettings::default();
41 cfg.set_host(host);
42 cfg.set_api_key(api_key);
43 Ok(Self {
44 inner,
45 cfg: Arc::new(cfg),
46 })
47 }
48
49 pub fn config(&self) -> &AppSearchSettings {
50 &self.cfg
51 }
52
53 pub fn host(&self) -> &str {
54 self.inner.get_host()
55 }
56
57 pub(crate) fn client(&self) -> &MeiliClient {
58 &self.inner
59 }
60
61 pub async fn is_healthy(&self) -> bool {
62 self.inner.is_healthy().await
63 }
64
65 pub async fn health(&self) -> SearchResult<Health> {
66 self.inner.health().await.map_err(SearchError::Meilisearch)
67 }
68
69 pub fn index(&self, uid: impl Into<String>) -> IndexHandle {
70 IndexHandle::new(
71 uid,
72 self.clone(),
73 Some(self.cfg.default_primary_key.clone()),
74 )
75 }
76
77 pub fn synced<T: crate::sync::Searchable>(&self) -> SyncedIndex<T> {
78 let uid = format!("{}.{}", self.cfg.index_prefix, T::INDEX_UID);
79 let mut handle = self.index(uid);
80 handle.set_primary_key(T::PRIMARY_KEY);
81 SyncedIndex::new(handle)
82 }
83
84 pub async fn delete_index(&self, uid: impl AsRef<str>) -> SearchResult<()> {
85 let _timer = metrics::Timer::start(metrics::OP_INDEX_DELETE);
86 let result = self.delete_index_inner(uid.as_ref()).await;
87 metrics::record_index_operation(metrics::OP_INDEX_DELETE, &result);
88 result
89 }
90
91 async fn delete_index_inner(&self, uid: &str) -> SearchResult<()> {
92 let task = self.inner.delete_index(uid).await?;
93 let task = self.wait_for_task(task).await?;
94 if task.is_success() || task_failed_code(&task) == Some(ErrorCode::IndexNotFound) {
95 Ok(())
96 } else {
97 Err(task_failure(&task))
98 }
99 }
100
101 pub async fn wait_for_task(&self, task: TaskInfo) -> SearchResult<Task> {
102 let timeout = Duration::from_millis(self.cfg.task_wait_timeout_ms);
103 let interval = Duration::from_millis(self.cfg.task_poll_interval_ms);
104 let task = self
105 .inner
106 .wait_for_task(task, Some(interval), Some(timeout))
107 .await?;
108 if let Task::Failed { content } = &task {
109 tracing::error!(
110 task_uid = task.get_uid(),
111 error_code = ?content.error.error_code,
112 error_message = %content.error.error_message,
113 "meilisearch task failed"
114 );
115 }
116 Ok(task)
117 }
118}
119
120pub struct SearchClientBuilder {
121 cfg: AppSearchSettings,
122}
123
124impl SearchClientBuilder {
125 pub fn from_config(cfg: AppSearchSettings) -> Self {
126 Self { cfg }
127 }
128
129 pub fn build(self) -> SearchResult<Option<SearchClient>> {
130 SearchClient::from_config(self.cfg)
131 }
132}
133
134pub(crate) fn task_failed_code(task: &Task) -> Option<ErrorCode> {
135 match task {
136 Task::Failed { content } => Some(content.error.error_code.clone()),
137 _ => None,
138 }
139}
140
141pub(crate) fn task_failure(task: &Task) -> SearchError {
142 match task {
143 Task::Failed { content } => SearchError::TaskFailed(format!(
144 "{:?}: {}",
145 content.error.error_code, content.error.error_message
146 )),
147 Task::Succeeded { .. } | Task::Enqueued { .. } | Task::Processing { .. } => {
148 SearchError::TaskFailed("task did not reach a successful state".to_string())
149 }
150 }
151}
152
153#[cfg(test)]
154mod tests {
155 use super::*;
156
157 #[test]
158 fn disabled_returns_none() {
159 let cfg = AppSearchSettings {
160 enabled: false,
161 ..Default::default()
162 };
163 assert!(
164 SearchClientBuilder::from_config(cfg)
165 .build()
166 .unwrap()
167 .is_none()
168 );
169 }
170
171 #[test]
172 fn invalid_config_rejected() {
173 let cfg = AppSearchSettings {
174 enabled: true,
175 host: "".into(),
176 ..Default::default()
177 };
178 assert!(SearchClientBuilder::from_config(cfg).build().is_err());
179 }
180
181 #[test]
182 fn construction_is_network_free() {
183 let client = SearchClient::new("http://127.0.0.1:7700", None::<String>).unwrap();
184 assert_eq!(client.host(), "http://127.0.0.1:7700");
185 }
186}
Added lib/search/error.rs lib/search/error.rs +59 −0 View lib/search/error.rs
@@ -0,0 +1,59 @@
1use thiserror::Error;
2
3pub type SearchResult<T> = Result<T, SearchError>;
4
5#[derive(Debug, Error)]
6pub enum SearchError {
7 #[error("meilisearch error: {0}")]
8 Meilisearch(#[from] meilisearch_sdk::errors::Error),
9 #[error("serialization error: {0}")]
10 Serde(#[from] serde_json::Error),
11 #[error("search is disabled")]
12 Disabled,
13 #[error("config invalid: {0}")]
14 Config(String),
15 #[error("index not found: {0}")]
16 IndexNotFound(String),
17 #[error("meilisearch task failed: {0}")]
18 TaskFailed(String),
19 #[error("invalid search query: {0}")]
20 InvalidQuery(String),
21 #[error("sync error: {0}")]
22 Sync(String),
23 #[error("queue error: {0}")]
24 Queue(#[from] queue::QueueError),
25}
26
27impl SearchError {
28 pub fn error_kind(&self) -> &'static str {
29 match self {
30 SearchError::Meilisearch(_) => "meilisearch",
31 SearchError::Serde(_) => "serde",
32 SearchError::Disabled => "disabled",
33 SearchError::Config(_) => "config",
34 SearchError::IndexNotFound(_) => "index_not_found",
35 SearchError::TaskFailed(_) => "task_failed",
36 SearchError::InvalidQuery(_) => "invalid_query",
37 SearchError::Sync(_) => "sync",
38 SearchError::Queue(_) => "queue",
39 }
40 }
41}
42
43#[cfg(test)]
44mod tests {
45 use super::*;
46
47 #[test]
48 fn error_kind_labels() {
49 assert_eq!(SearchError::Disabled.error_kind(), "disabled");
50 assert_eq!(
51 SearchError::InvalidQuery("x".into()).error_kind(),
52 "invalid_query"
53 );
54 assert_eq!(
55 SearchError::TaskFailed("x".into()).error_kind(),
56 "task_failed"
57 );
58 }
59}
Added lib/search/index.rs lib/search/index.rs +394 −0 View lib/search/index.rs
@@ -0,0 +1,394 @@
1use meilisearch_sdk::errors::ErrorCode;
2use meilisearch_sdk::indexes::Index as SdkIndex;
3use meilisearch_sdk::search::{SearchResults, Selectors};
4use meilisearch_sdk::settings::PaginationSetting;
5use meilisearch_sdk::task_info::TaskInfo;
6use serde::Serialize;
7use serde::de::DeserializeOwned;
8
9use crate::client::task_failure;
10use crate::error::{SearchError, SearchResult};
11use crate::metrics;
12use crate::query::{SearchHit, SearchRequest, SearchResponse};
13
14#[derive(Debug, Clone, Default, PartialEq, Eq)]
15pub struct IndexSettings {
16 pub primary_key: Option<String>,
17 pub searchable_attributes: Vec<String>,
18 pub filterable_attributes: Vec<String>,
19 pub sortable_attributes: Vec<String>,
20 pub displayed_attributes: Vec<String>,
21 pub ranking_rules: Vec<String>,
22 pub distinct_attribute: Option<String>,
23 pub max_total_hits: Option<usize>,
24}
25
26impl IndexSettings {
27 pub fn with_primary_key(mut self, pk: impl Into<String>) -> Self {
28 self.primary_key = Some(pk.into());
29 self
30 }
31
32 pub fn with_searchable_attributes(
33 mut self,
34 attrs: impl IntoIterator<Item = impl Into<String>>,
35 ) -> Self {
36 self.searchable_attributes = attrs.into_iter().map(Into::into).collect();
37 self
38 }
39
40 pub fn with_filterable_attributes(
41 mut self,
42 attrs: impl IntoIterator<Item = impl Into<String>>,
43 ) -> Self {
44 self.filterable_attributes = attrs.into_iter().map(Into::into).collect();
45 self
46 }
47
48 pub fn with_sortable_attributes(
49 mut self,
50 attrs: impl IntoIterator<Item = impl Into<String>>,
51 ) -> Self {
52 self.sortable_attributes = attrs.into_iter().map(Into::into).collect();
53 self
54 }
55
56 pub fn with_displayed_attributes(
57 mut self,
58 attrs: impl IntoIterator<Item = impl Into<String>>,
59 ) -> Self {
60 self.displayed_attributes = attrs.into_iter().map(Into::into).collect();
61 self
62 }
63
64 pub fn with_ranking_rules(
65 mut self,
66 rules: impl IntoIterator<Item = impl Into<String>>,
67 ) -> Self {
68 self.ranking_rules = rules.into_iter().map(Into::into).collect();
69 self
70 }
71
72 pub fn with_distinct_attribute(mut self, attr: impl Into<String>) -> Self {
73 self.distinct_attribute = Some(attr.into());
74 self
75 }
76
77 pub fn with_max_total_hits(mut self, max: usize) -> Self {
78 self.max_total_hits = Some(max);
79 self
80 }
81}
82
83#[derive(Clone)]
84pub struct IndexHandle {
85 uid: String,
86 client: crate::client::SearchClient,
87 primary_key: Option<String>,
88}
89
90impl IndexHandle {
91 pub(crate) fn new(
92 uid: impl Into<String>,
93 client: crate::client::SearchClient,
94 primary_key: Option<String>,
95 ) -> Self {
96 Self {
97 uid: uid.into(),
98 client,
99 primary_key,
100 }
101 }
102
103 pub fn uid(&self) -> &str {
104 &self.uid
105 }
106
107 pub fn client(&self) -> &crate::client::SearchClient {
108 &self.client
109 }
110
111 pub fn primary_key(&self) -> Option<&str> {
112 self.primary_key.as_deref()
113 }
114
115 pub fn set_primary_key(&mut self, pk: impl Into<String>) {
116 self.primary_key = Some(pk.into());
117 }
118
119 pub fn batch_size(&self) -> usize {
120 self.client.config().batch_size
121 }
122
123 pub(crate) fn sdk(&self) -> SdkIndex {
124 self.client.client().index(&self.uid)
125 }
126
127 pub async fn ensure_index(&self, settings: &IndexSettings) -> SearchResult<()> {
128 let _timer = metrics::Timer::start(metrics::OP_INDEX_CREATE);
129 let result = self.ensure_index_inner(settings).await;
130 metrics::record_index_operation(metrics::OP_INDEX_CREATE, &result);
131 result
132 }
133
134 pub(crate) async fn ensure_index_inner(&self, settings: &IndexSettings) -> SearchResult<()> {
135 let client = self.client.client();
136 let exists = match client.get_index(&self.uid).await {
137 Ok(_) => true,
138 Err(e) if is_index_not_found(&e) => false,
139 Err(e) => return Err(SearchError::Meilisearch(e)),
140 };
141 if !exists {
142 let task = client
143 .create_index(&self.uid, settings.primary_key.as_deref())
144 .await?;
145 self.await_task(task).await?;
146 }
147 self.apply_settings(settings).await
148 }
149
150 pub async fn apply_settings(&self, settings: &IndexSettings) -> SearchResult<()> {
151 let _timer = metrics::Timer::start(metrics::OP_INDEX_SETTINGS);
152 let index = self.sdk();
153 let mut tasks = Vec::with_capacity(7);
154 if !settings.searchable_attributes.is_empty() {
155 tasks.push(
156 index
157 .set_searchable_attributes(&settings.searchable_attributes)
158 .await?,
159 );
160 }
161 if !settings.filterable_attributes.is_empty() {
162 tasks.push(
163 index
164 .set_filterable_attributes(&settings.filterable_attributes)
165 .await?,
166 );
167 }
168 if !settings.sortable_attributes.is_empty() {
169 tasks.push(
170 index
171 .set_sortable_attributes(&settings.sortable_attributes)
172 .await?,
173 );
174 }
175 if !settings.displayed_attributes.is_empty() {
176 tasks.push(
177 index
178 .set_displayed_attributes(&settings.displayed_attributes)
179 .await?,
180 );
181 }
182 if !settings.ranking_rules.is_empty() {
183 tasks.push(index.set_ranking_rules(&settings.ranking_rules).await?);
184 }
185 if let Some(distinct) = &settings.distinct_attribute {
186 tasks.push(index.set_distinct_attribute(distinct).await?);
187 }
188 if let Some(max_total_hits) = settings.max_total_hits {
189 tasks.push(
190 index
191 .set_pagination(PaginationSetting { max_total_hits })
192 .await?,
193 );
194 }
195 for task in tasks {
196 self.await_task(task).await?;
197 }
198 metrics::record_index_operation(metrics::OP_INDEX_SETTINGS, &Ok(()));
199 Ok(())
200 }
201
202 pub async fn delete(&self) -> SearchResult<()> {
203 let _timer = metrics::Timer::start(metrics::OP_INDEX_DELETE);
204 let task = self.sdk().delete().await?;
205 let result = self.await_task(task).await;
206 metrics::record_index_operation(metrics::OP_INDEX_DELETE, &result);
207 result
208 }
209
210 pub async fn add_documents<T: Serialize + Send + Sync>(
211 &self,
212 documents: &[T],
213 ) -> SearchResult<()> {
214 if documents.is_empty() {
215 return Ok(());
216 }
217 let task = self
218 .sdk()
219 .add_documents(documents, self.primary_key.as_deref())
220 .await?;
221 self.await_task(task).await?;
222 metrics::record_documents(&self.uid, metrics::OP_DOCS_ADD, documents.len());
223 Ok(())
224 }
225
226 pub async fn update_documents<T: Serialize + Send + Sync>(
227 &self,
228 documents: &[T],
229 ) -> SearchResult<()> {
230 if documents.is_empty() {
231 return Ok(());
232 }
233 let task = self
234 .sdk()
235 .add_or_update(documents, self.primary_key.as_deref())
236 .await?;
237 self.await_task(task).await?;
238 metrics::record_documents(&self.uid, metrics::OP_DOCS_UPDATE, documents.len());
239 Ok(())
240 }
241
242 pub async fn delete_documents(&self, ids: &[String]) -> SearchResult<()> {
243 if ids.is_empty() {
244 return Ok(());
245 }
246 let task = self.sdk().delete_documents(ids).await?;
247 self.await_task(task).await?;
248 metrics::record_documents(&self.uid, metrics::OP_DOCS_DELETE, ids.len());
249 Ok(())
250 }
251
252 pub async fn delete_document(&self, id: &str) -> SearchResult<()> {
253 let task = self.sdk().delete_document(id).await?;
254 self.await_task(task).await?;
255 metrics::record_documents(&self.uid, metrics::OP_DOCS_DELETE, 1);
256 Ok(())
257 }
258
259 pub async fn clear(&self) -> SearchResult<()> {
260 let task = self.sdk().delete_all_documents().await?;
261 self.await_task(task).await?;
262 metrics::record_documents(&self.uid, metrics::OP_DOCS_CLEAR, 0);
263 Ok(())
264 }
265
266 pub async fn search<T: DeserializeOwned + Send + Sync + 'static>(
267 &self,
268 request: &SearchRequest,
269 ) -> SearchResult<SearchResponse<T>> {
270 request.validate()?;
271 let _timer = metrics::Timer::start(metrics::OP_SEARCH);
272 let result = self.execute_search::<T>(request).await;
273 metrics::record_request(metrics::OP_SEARCH, &result);
274 result
275 }
276
277 pub async fn search_raw(
278 &self,
279 request: &SearchRequest,
280 ) -> SearchResult<SearchResponse<serde_json::Value>> {
281 self.search::<serde_json::Value>(request).await
282 }
283
284 async fn execute_search<T: DeserializeOwned + Send + Sync + 'static>(
285 &self,
286 request: &SearchRequest,
287 ) -> SearchResult<SearchResponse<T>> {
288 let sort_refs = request
289 .sort
290 .as_ref()
291 .map(|v| v.iter().map(String::as_str).collect::<Vec<_>>());
292 let retrieve_refs = request
293 .attributes_to_retrieve
294 .as_ref()
295 .map(|v| v.iter().map(String::as_str).collect::<Vec<_>>());
296 let highlight_refs = request
297 .attributes_to_highlight
298 .as_ref()
299 .map(|v| v.iter().map(String::as_str).collect::<Vec<_>>());
300 let facet_refs = request
301 .facets
302 .as_ref()
303 .map(|v| v.iter().map(String::as_str).collect::<Vec<_>>());
304
305 let index = self.sdk();
306 let mut q = index.search();
307 q.with_query(&request.q);
308 q.with_limit(request.limit);
309 q.with_offset(request.offset);
310 if let Some(filter) = &request.filter {
311 q.with_filter(filter);
312 }
313 if let Some(refs) = &sort_refs {
314 q.with_sort(refs);
315 }
316 if let Some(refs) = &retrieve_refs {
317 q.with_attributes_to_retrieve(Selectors::Some(refs));
318 }
319 if let Some(refs) = &highlight_refs {
320 q.with_attributes_to_highlight(Selectors::Some(refs));
321 }
322 if let Some(refs) = &facet_refs {
323 q.with_facets(Selectors::Some(refs));
324 }
325 if let Some(distinct) = &request.distinct {
326 q.with_distinct(distinct);
327 }
328 if request.show_ranking_score {
329 q.with_show_ranking_score(true);
330 }
331 let built = q.build();
332 let results = built.execute::<T>().await?;
333 Ok(SearchResponse::from_sdk(&self.uid, results))
334 }
335
336 async fn await_task(&self, task: TaskInfo) -> SearchResult<()> {
337 let task = self.client.wait_for_task(task).await?;
338 if task.is_success() {
339 Ok(())
340 } else {
341 Err(task_failure(&task))
342 }
343 }
344}
345
346pub(crate) fn is_index_not_found(e: &meilisearch_sdk::errors::Error) -> bool {
347 matches!(
348 e,
349 meilisearch_sdk::errors::Error::Meilisearch(me) if me.error_code == ErrorCode::IndexNotFound
350 )
351}
352
353impl<T> SearchResponse<T> {
354 pub(crate) fn from_sdk(index_uid: &str, r: SearchResults<T>) -> Self {
355 let hits = r
356 .hits
357 .into_iter()
358 .map(|h| SearchHit {
359 document: h.result,
360 ranking_score: h.ranking_score,
361 formatted: h.formatted_result,
362 })
363 .collect();
364 Self {
365 hits,
366 offset: r.offset.unwrap_or(0),
367 limit: r.limit.unwrap_or(0),
368 estimated_total_hits: r.estimated_total_hits.unwrap_or(0),
369 total_hits: r.total_hits,
370 total_pages: r.total_pages,
371 processing_time_ms: r.processing_time_ms,
372 query: r.query,
373 index_uid: r.index_uid.or_else(|| Some(index_uid.to_string())),
374 }
375 }
376}
377
378#[cfg(test)]
379mod tests {
380 use super::*;
381
382 #[test]
383 fn settings_builder() {
384 let s = IndexSettings::default()
385 .with_primary_key("uid")
386 .with_searchable_attributes(["name", "description"])
387 .with_filterable_attributes(["space_id"])
388 .with_sortable_attributes(["updated_at"])
389 .with_max_total_hits(10_000);
390 assert_eq!(s.primary_key.as_deref(), Some("uid"));
391 assert_eq!(s.searchable_attributes, vec!["name", "description"]);
392 assert_eq!(s.max_total_hits, Some(10_000));
393 }
394}
Added lib/search/lib.rs lib/search/lib.rs +13 −0 View lib/search/lib.rs
@@ -0,0 +1,13 @@
1//! .expect("search enabled");
2pub mod client;
3pub mod error;
4pub mod index;
5pub mod metrics;
6pub mod query;
7pub mod sync;
8
9pub use client::{SearchClient, SearchClientBuilder};
10pub use error::{SearchError, SearchResult};
11pub use index::{IndexHandle, IndexSettings};
12pub use query::{SearchHit, SearchRequest, SearchResponse};
13pub use sync::{Searchable, SyncBridge, SyncMessage, SyncSummary, SyncedIndex};
Added lib/search/metrics.rs lib/search/metrics.rs +90 −0 View lib/search/metrics.rs
@@ -0,0 +1,90 @@
1use std::time::Instant;
2
3use metrics::{counter, describe_counter, describe_histogram, histogram};
4
5use crate::error::SearchError;
6
7pub const OP_SEARCH: &str = "search";
8pub const OP_INDEX_CREATE: &str = "index_create";
9pub const OP_INDEX_DELETE: &str = "index_delete";
10pub const OP_INDEX_SETTINGS: &str = "index_settings";
11pub const OP_DOCS_ADD: &str = "docs_add";
12pub const OP_DOCS_UPDATE: &str = "docs_update";
13pub const OP_DOCS_DELETE: &str = "docs_delete";
14pub const OP_DOCS_CLEAR: &str = "docs_clear";
15
16pub const OUTCOME_SUCCESS: &str = "success";
17
18static DESCRIBE_ONCE: std::sync::Once = std::sync::Once::new();
19
20pub fn ensure_described() {
21 DESCRIBE_ONCE.call_once(describe_all);
22}
23
24fn describe_all() {
25 describe_counter!(
26 "search_requests_total",
27 "Total search / mutation requests labeled by operation and outcome"
28 );
29 describe_histogram!(
30 "search_request_duration_seconds",
31 metrics::Unit::Seconds,
32 "Search and index operation latency"
33 );
34 describe_counter!(
35 "search_documents_total",
36 "Documents written to Meilisearch labeled by index and operation"
37 );
38}
39
40pub struct Timer {
41 pub op: &'static str,
42 pub start: Instant,
43}
44
45impl Timer {
46 pub fn start(op: &'static str) -> Self {
47 ensure_described();
48 Self {
49 op,
50 start: Instant::now(),
51 }
52 }
53}
54
55impl Drop for Timer {
56 fn drop(&mut self) {
57 let elapsed = self.start.elapsed().as_secs_f64();
58 histogram!(
59 "search_request_duration_seconds",
60 "operation" => self.op.to_string()
61 )
62 .record(elapsed);
63 }
64}
65
66pub fn record_request<T>(op: &'static str, result: &Result<T, SearchError>) {
67 let outcome = match result {
68 Ok(_) => OUTCOME_SUCCESS,
69 Err(e) => e.error_kind(),
70 };
71 counter!(
72 "search_requests_total",
73 "operation" => op.to_string(),
74 "outcome" => outcome.to_string()
75 )
76 .increment(1);
77}
78
79pub fn record_documents(index: &str, op: &'static str, count: usize) {
80 counter!(
81 "search_documents_total",
82 "index" => index.to_string(),
83 "operation" => op.to_string()
84 )
85 .increment(count as u64);
86}
87
88pub fn record_index_operation<T>(op: &'static str, result: &Result<T, SearchError>) {
89 record_request(op, result);
90}
Added lib/search/query.rs lib/search/query.rs +200 −0 View lib/search/query.rs
@@ -0,0 +1,200 @@
1use serde::{Deserialize, Serialize};
2
3use crate::error::{SearchError, SearchResult};
4
5pub const MAX_LIMIT: usize = 1_000;
6
7fn default_limit() -> usize {
8 20
9}
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub struct SearchRequest {
13 #[serde(default)]
14 pub q: String,
15 #[serde(default = "default_limit")]
16 pub limit: usize,
17 #[serde(default)]
18 pub offset: usize,
19 #[serde(default, skip_serializing_if = "Option::is_none")]
20 pub filter: Option<String>,
21 #[serde(default, skip_serializing_if = "Option::is_none")]
22 pub sort: Option<Vec<String>>,
23 #[serde(default, skip_serializing_if = "Option::is_none")]
24 pub attributes_to_retrieve: Option<Vec<String>>,
25 #[serde(default, skip_serializing_if = "Option::is_none")]
26 pub attributes_to_highlight: Option<Vec<String>>,
27 #[serde(default, skip_serializing_if = "Option::is_none")]
28 pub distinct: Option<String>,
29 #[serde(default, skip_serializing_if = "Option::is_none")]
30 pub facets: Option<Vec<String>>,
31 #[serde(default)]
32 pub show_ranking_score: bool,
33}
34
35impl Default for SearchRequest {
36 fn default() -> Self {
37 Self {
38 q: String::new(),
39 limit: default_limit(),
40 offset: 0,
41 filter: None,
42 sort: None,
43 attributes_to_retrieve: None,
44 attributes_to_highlight: None,
45 distinct: None,
46 facets: None,
47 show_ranking_score: false,
48 }
49 }
50}
51
52impl SearchRequest {
53 pub fn new(q: impl Into<String>) -> Self {
54 Self {
55 q: q.into(),
56 ..Self::default()
57 }
58 }
59
60 pub fn with_limit(mut self, limit: usize) -> Self {
61 self.limit = limit;
62 self
63 }
64
65 pub fn with_offset(mut self, offset: usize) -> Self {
66 self.offset = offset;
67 self
68 }
69
70 pub fn with_filter(mut self, filter: impl Into<String>) -> Self {
71 self.filter = Some(filter.into());
72 self
73 }
74
75 pub fn with_sort(mut self, sort: impl IntoIterator<Item = impl Into<String>>) -> Self {
76 self.sort = Some(sort.into_iter().map(Into::into).collect());
77 self
78 }
79
80 pub fn with_attributes_to_retrieve(
81 mut self,
82 attrs: impl IntoIterator<Item = impl Into<String>>,
83 ) -> Self {
84 self.attributes_to_retrieve = Some(attrs.into_iter().map(Into::into).collect());
85 self
86 }
87
88 pub fn with_attributes_to_highlight(
89 mut self,
90 attrs: impl IntoIterator<Item = impl Into<String>>,
91 ) -> Self {
92 self.attributes_to_highlight = Some(attrs.into_iter().map(Into::into).collect());
93 self
94 }
95
96 pub fn with_distinct(mut self, attr: impl Into<String>) -> Self {
97 self.distinct = Some(attr.into());
98 self
99 }
100
101 pub fn with_facets(mut self, facets: impl IntoIterator<Item = impl Into<String>>) -> Self {
102 self.facets = Some(facets.into_iter().map(Into::into).collect());
103 self
104 }
105
106 pub fn with_show_ranking_score(mut self, show: bool) -> Self {
107 self.show_ranking_score = show;
108 self
109 }
110
111 pub fn validate(&self) -> SearchResult<()> {
112 if self.limit == 0 || self.limit > MAX_LIMIT {
113 return Err(SearchError::InvalidQuery(format!(
114 "limit must be in 1..={MAX_LIMIT}, got {}",
115 self.limit
116 )));
117 }
118 Ok(())
119 }
120}
121
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct SearchHit<T> {
124 #[serde(flatten)]
125 pub document: T,
126 #[serde(default, skip_serializing_if = "Option::is_none")]
127 pub ranking_score: Option<f64>,
128 #[serde(
129 default,
130 rename = "_formatted",
131 skip_serializing_if = "Option::is_none"
132 )]
133 pub formatted: Option<serde_json::Map<String, serde_json::Value>>,
134}
135
136#[derive(Debug, Clone, Serialize, Deserialize)]
137pub struct SearchResponse<T> {
138 pub hits: Vec<SearchHit<T>>,
139 pub offset: usize,
140 pub limit: usize,
141 pub estimated_total_hits: usize,
142 pub total_hits: Option<usize>,
143 pub total_pages: Option<usize>,
144 pub processing_time_ms: usize,
145 pub query: String,
146 pub index_uid: Option<String>,
147}
148
149#[cfg(test)]
150mod tests {
151 use super::*;
152
153 #[test]
154 fn default_limit_is_twenty() {
155 assert_eq!(SearchRequest::new("x").limit, 20);
156 }
157
158 #[test]
159 fn builder_chains() {
160 let req = SearchRequest::new("rust")
161 .with_limit(10)
162 .with_offset(5)
163 .with_filter("space_id = \"s1\"")
164 .with_sort(["updated_at:desc"])
165 .with_show_ranking_score(true);
166 assert_eq!(req.limit, 10);
167 assert_eq!(req.offset, 5);
168 assert_eq!(req.filter.as_deref(), Some("space_id = \"s1\""));
169 assert_eq!(
170 req.sort.as_deref(),
171 Some(&["updated_at:desc".to_string()][..])
172 );
173 assert!(req.show_ranking_score);
174 }
175
176 #[test]
177 fn validate_rejects_zero_limit() {
178 assert!(SearchRequest::new("x").with_limit(0).validate().is_err());
179 }
180
181 #[test]
182 fn validate_rejects_excessive_limit() {
183 assert!(
184 SearchRequest::new("x")
185 .with_limit(MAX_LIMIT + 1)
186 .validate()
187 .is_err()
188 );
189 }
190
191 #[test]
192 fn serde_roundtrip() {
193 let req = SearchRequest::new("rust")
194 .with_limit(5)
195 .with_filter("a = 1");
196 let json = serde_json::to_string(&req).unwrap();
197 let back: SearchRequest = serde_json::from_str(&json).unwrap();
198 assert_eq!(req, back);
199 }
200}
Added lib/search/sync.rs lib/search/sync.rs +295 −0 View lib/search/sync.rs
@@ -0,0 +1,295 @@
1use std::marker::PhantomData;
2
3use queue::consumer::SubscribeHandle;
4use queue::typed::TypedMessage;
5use serde::de::DeserializeOwned;
6use serde::{Deserialize, Serialize};
7
8use crate::error::{SearchError, SearchResult};
9use crate::index::{IndexHandle, IndexSettings};
10use crate::metrics;
11use crate::query::{SearchRequest, SearchResponse};
12
13pub trait Searchable: Send + Sync + 'static {
14 const INDEX_UID: &'static str;
15 const PRIMARY_KEY: &'static str;
16
17 fn index_settings() -> IndexSettings {
18 IndexSettings::default()
19 }
20
21 fn document_id(&self) -> String;
22}
23
24#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
25#[serde(tag = "op", content = "payload", rename_all = "snake_case")]
26pub enum SyncMessage<T> {
27 Upsert { document: T },
28 Update { document: T },
29 Delete { id: String },
30 Clear {},
31}
32
33#[derive(Debug, Clone, Default)]
34pub struct SyncSummary {
35 pub index: String,
36 pub operation: &'static str,
37 pub documents: usize,
38}
39
40impl SyncSummary {
41 pub fn new(index: &str, operation: &'static str, documents: usize) -> Self {
42 Self {
43 index: index.to_string(),
44 operation,
45 documents,
46 }
47 }
48}
49
50pub struct SyncedIndex<T> {
51 index: IndexHandle,
52 _marker: PhantomData<T>,
53}
54
55impl<T> Clone for SyncedIndex<T> {
56 fn clone(&self) -> Self {
57 Self {
58 index: self.index.clone(),
59 _marker: PhantomData,
60 }
61 }
62}
63
64impl<T> SyncedIndex<T>
65where
66 T: Searchable,
67{
68 pub(crate) fn new(index: IndexHandle) -> Self {
69 Self {
70 index,
71 _marker: PhantomData,
72 }
73 }
74
75 pub fn uid(&self) -> &str {
76 self.index.uid()
77 }
78
79 pub fn handle(&self) -> &IndexHandle {
80 &self.index
81 }
82
83 pub async fn ensure_index(&self) -> SearchResult<()> {
84 self.index.ensure_index(&T::index_settings()).await
85 }
86}
87
88impl<T> SyncedIndex<T>
89where
90 T: Searchable + Clone + Serialize + DeserializeOwned + Send + Sync,
91{
92 pub fn bridge<'a>(&self, queue: &'a queue::Queue) -> SyncBridge<'a, T> {
93 SyncBridge::new(queue, self.clone())
94 }
95}
96
97impl<T> SyncedIndex<T>
98where
99 T: Searchable + Serialize,
100{
101 pub async fn upsert(&self, documents: &[T]) -> SearchResult<SyncSummary> {
102 let _timer = metrics::Timer::start(metrics::OP_DOCS_ADD);
103 let batch = self.index.batch_size();
104 for chunk in documents.chunks(batch) {
105 self.index.add_documents(chunk).await?;
106 }
107 Ok(SyncSummary::new(
108 self.uid(),
109 metrics::OP_DOCS_ADD,
110 documents.len(),
111 ))
112 }
113
114 pub async fn update(&self, documents: &[T]) -> SearchResult<SyncSummary> {
115 let _timer = metrics::Timer::start(metrics::OP_DOCS_UPDATE);
116 let batch = self.index.batch_size();
117 for chunk in documents.chunks(batch) {
118 self.index.update_documents(chunk).await?;
119 }
120 Ok(SyncSummary::new(
121 self.uid(),
122 metrics::OP_DOCS_UPDATE,
123 documents.len(),
124 ))
125 }
126
127 pub async fn delete(&self, ids: &[String]) -> SearchResult<SyncSummary> {
128 self.index.delete_documents(ids).await?;
129 Ok(SyncSummary::new(
130 self.uid(),
131 metrics::OP_DOCS_DELETE,
132 ids.len(),
133 ))
134 }
135
136 pub async fn clear(&self) -> SearchResult<SyncSummary> {
137 self.index.clear().await?;
138 Ok(SyncSummary::new(self.uid(), metrics::OP_DOCS_CLEAR, 0))
139 }
140}
141
142impl<T> SyncedIndex<T>
143where
144 T: Searchable + DeserializeOwned,
145{
146 pub async fn search(&self, request: &SearchRequest) -> SearchResult<SearchResponse<T>> {
147 self.index.search(request).await
148 }
149}
150
151pub struct SyncBridge<'a, T> {
152 queue: &'a queue::Queue,
153 index: SyncedIndex<T>,
154}
155
156impl<'a, T> SyncBridge<'a, T>
157where
158 T: Searchable + Clone + Serialize + DeserializeOwned + Send + Sync,
159{
160 pub(crate) fn new(queue: &'a queue::Queue, index: SyncedIndex<T>) -> Self {
161 Self { queue, index }
162 }
163
164 pub fn index(&self) -> &SyncedIndex<T> {
165 &self.index
166 }
167
168 pub async fn ensure_index(&self) -> SearchResult<()> {
169 self.index.ensure_index().await
170 }
171
172 pub async fn publish_upsert(&self, document: &T) -> SearchResult<()> {
173 self.publish(&SyncMessage::Upsert {
174 document: document.clone(),
175 })
176 .await
177 }
178
179 pub async fn publish_update(&self, document: &T) -> SearchResult<()> {
180 self.publish(&SyncMessage::Update {
181 document: document.clone(),
182 })
183 .await
184 }
185
186 pub async fn publish_delete(&self, id: impl Into<String>) -> SearchResult<()> {
187 self.publish(&SyncMessage::Delete { id: id.into() }).await
188 }
189
190 pub async fn publish_clear(&self) -> SearchResult<()> {
191 self.publish(&SyncMessage::Clear {}).await
192 }
193
194 async fn publish(&self, message: &SyncMessage<T>) -> SearchResult<()> {
195 self.queue
196 .typed::<SyncMessage<T>>()
197 .publish(message)
198 .await
199 .map_err(SearchError::Queue)
200 }
201
202 pub async fn spawn_consumer(
203 &self,
204 mode: queue::SubscriptionMode,
205 ) -> Result<SubscribeHandle, SearchError> {
206 self.index.ensure_index().await?;
207 let index = self.index.clone();
208 let handle = self
209 .queue
210 .typed::<SyncMessage<T>>()
211 .subscribe(mode, move |msg: TypedMessage<SyncMessage<T>>| {
212 let index = index.clone();
213 async move {
214 let result = match msg.value {
215 SyncMessage::Upsert { document } => index.upsert(&[document]).await,
216 SyncMessage::Update { document } => index.update(&[document]).await,
217 SyncMessage::Delete { id } => index.delete(&[id]).await,
218 SyncMessage::Clear {} => index.clear().await,
219 };
220 match result {
221 Ok(summary) => {
222 tracing::debug!(
223 index = %summary.index,
224 operation = summary.operation,
225 documents = summary.documents,
226 "search sync applied"
227 );
228 Ok(())
229 }
230 Err(e) => {
231 tracing::warn!(
232 index = %index.uid(),
233 topic = %msg.meta.topic,
234 error = %e,
235 "search sync message failed"
236 );
237 Err(queue::QueueError::Handler(e.to_string()))
238 }
239 }
240 }
241 })
242 .map_err(SearchError::Queue)?;
243 Ok(handle)
244 }
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250 use crate::client::SearchClient;
251
252 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
253 struct Dummy {
254 id: String,
255 name: String,
256 }
257
258 impl Searchable for Dummy {
259 const INDEX_UID: &'static str = "dummy";
260 const PRIMARY_KEY: &'static str = "id";
261 fn document_id(&self) -> String {
262 self.id.clone()
263 }
264 }
265
266 #[test]
267 fn sync_message_roundtrip() {
268 let msg = SyncMessage::Upsert {
269 document: Dummy {
270 id: "1".into(),
271 name: "x".into(),
272 },
273 };
274 let json = serde_json::to_string(&msg).unwrap();
275 assert!(json.contains("\"op\":\"upsert\""));
276 let back: SyncMessage<Dummy> = serde_json::from_str(&json).unwrap();
277 assert_eq!(msg, back);
278 }
279
280 #[test]
281 fn delete_message_roundtrip() {
282 let msg = SyncMessage::<Dummy>::Delete { id: "42".into() };
283 let json = serde_json::to_string(&msg).unwrap();
284 assert!(json.contains("\"op\":\"delete\""));
285 let back: SyncMessage<Dummy> = serde_json::from_str(&json).unwrap();
286 assert_eq!(msg, back);
287 }
288
289 #[test]
290 fn synced_uid_respects_prefix() {
291 let client = SearchClient::new("http://127.0.0.1:7700", None::<String>).unwrap();
292 let synced = client.synced::<Dummy>();
293 assert_eq!(synced.uid(), "untitled.dummy");
294 }
295}