Skip to main content

sift_query_sqlx/execution/
search.rs

1use super::*;
2
3impl PgSiftQuery {
4    /// Executes a cursor page with one bounded list query. The query must use
5    /// the cursor paging mode (`cursor=` for the first page); no COUNT or
6    /// summary query is issued.
7    pub async fn search_cursor<E>(
8        &self,
9        query: impl QueryInput<E>,
10        scope: &ServerScope<E>,
11    ) -> Result<CursorPage<E>, PgSearchError>
12    where
13        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
14    {
15        self.observe::<E, _>(
16            "search_cursor",
17            async {
18                let mut plan =
19                    self.resolve_plan(query, scope, QueryTarget::TypedPage, "search_cursor")?;
20                let size = cursor_page_size(&mut plan)?;
21                if plan.force_empty || size == 0 {
22                    return Ok(CursorPage::empty());
23                }
24                let partition = self.extensions.resolve_partition(&plan, "search_cursor")?;
25                let rendered = prepare(self.render::<E>(
26                    &plan,
27                    "search_cursor",
28                    "list",
29                    partition.as_deref(),
30                    RenderRequest::List,
31                )?)?;
32                let deadline = self.deadline::<E>();
33                let connection = self.acquire(deadline).await?;
34                let (_, mut data) =
35                    fetch_typed::<E>(connection, rendered, deadline, TimeoutStage::ListQuery)
36                        .await?;
37                let has_more = data.len() > size as usize;
38                if has_more {
39                    data.truncate(size as usize);
40                }
41                let next_cursor = if has_more {
42                    self.entity_cursor(&plan, data.last())?
43                } else {
44                    None
45                };
46                Ok(CursorPage {
47                    data,
48                    next_cursor,
49                    has_more,
50                })
51            },
52            |page| page.data.len(),
53        )
54        .await
55    }
56
57    /// Dynamic map equivalent of [`Self::search_cursor`].
58    pub async fn search_map_cursor<E>(
59        &self,
60        query: impl QueryInput<E>,
61        scope: &ServerScope<E>,
62    ) -> Result<CursorPage<RowMap>, PgSearchError>
63    where
64        E: SearchEntity,
65    {
66        self.observe::<E, _>(
67            "search_map_cursor",
68            async {
69                let mut plan =
70                    self.resolve_plan(query, scope, QueryTarget::MapPage, "search_map_cursor")?;
71                let size = cursor_page_size(&mut plan)?;
72                if plan.force_empty || size == 0 {
73                    return Ok(CursorPage::empty());
74                }
75                let partition = self
76                    .extensions
77                    .resolve_partition(&plan, "search_map_cursor")?;
78                let rendered = prepare(self.render::<E>(
79                    &plan,
80                    "search_map_cursor",
81                    "list",
82                    partition.as_deref(),
83                    RenderRequest::List,
84                )?)?;
85                let deadline = self.deadline::<E>();
86                let connection = self.acquire(deadline).await?;
87                let (_, mut data) = fetch_maps(connection, rendered, &plan, deadline).await?;
88                let has_more = data.len() > size as usize;
89                if has_more {
90                    data.truncate(size as usize);
91                }
92                // Derive the token from the last database row before any
93                // formatting or result filter can remove it. Otherwise a
94                // post-filtered page could repeat the same underlying rows.
95                let next_cursor = if has_more {
96                    self.map_cursor(&plan, data.last())?
97                } else {
98                    None
99                };
100                self.format_map_rows::<E>(&mut data)?;
101                self.extensions
102                    .filter_rows::<E>("search_map_cursor", &mut data)?;
103                let page = self
104                    .extensions
105                    .enrich_page(
106                        E::META.entity_name,
107                        "search_map_cursor",
108                        Page {
109                            data,
110                            total_count: -1,
111                            summaries: Vec::new(),
112                            next_cursor: next_cursor.clone(),
113                        },
114                    )
115                    .await?;
116                Ok(CursorPage {
117                    data: page.data,
118                    next_cursor,
119                    has_more,
120                })
121            },
122            |page| page.data.len(),
123        )
124        .await
125    }
126
127    /// Executes a query after deriving its mandatory server scope from a
128    /// trusted application context.
129    pub async fn search_authorized<E, C, P>(
130        &self,
131        query: impl QueryInput<E>,
132        context: &C,
133        provider: &P,
134    ) -> Result<Page<E>, PgSearchError>
135    where
136        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
137        P: sift_query_core::ScopeProvider<C, E>,
138    {
139        let scope = provider.scope(context)?;
140        self.search(query, &scope).await
141    }
142
143    /// Dynamic map equivalent of [`Self::search_authorized`].
144    pub async fn search_map_authorized<E, C, P>(
145        &self,
146        query: impl QueryInput<E>,
147        context: &C,
148        provider: &P,
149    ) -> Result<Page<RowMap>, PgSearchError>
150    where
151        E: SearchEntity,
152        P: sift_query_core::ScopeProvider<C, E>,
153    {
154        let scope = provider.scope(context)?;
155        self.search_map(query, &scope).await
156    }
157
158    /// Count equivalent of [`Self::search_authorized`].
159    pub async fn count_authorized<E, C, P>(
160        &self,
161        query: impl QueryInput<E>,
162        context: &C,
163        provider: &P,
164    ) -> Result<i64, PgSearchError>
165    where
166        E: SearchEntity,
167        P: sift_query_core::ScopeProvider<C, E>,
168    {
169        let scope = provider.scope(context)?;
170        self.count(query, &scope).await
171    }
172
173    /// Runs typed page enrichment only when explicitly requested by the caller.
174    pub async fn search_enriched<E, R>(
175        &self,
176        query: impl QueryInput<E>,
177        scope: &ServerScope<E>,
178        enricher: &R,
179    ) -> Result<Page<E>, PgSearchError>
180    where
181        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
182        R: crate::TypedPageEnricher<E>,
183    {
184        self.observe::<E, _>(
185            "search_enriched",
186            async {
187                let mut page = self.search(query, scope).await?;
188                let expected_rows = page.data.len();
189                page.data = tokio::time::timeout(
190                    self.execution.timeout,
191                    enricher.enrich(
192                        page.data,
193                        crate::EnrichContext {
194                            entity_name: E::META.entity_name,
195                            operation: "search_enriched",
196                        },
197                    ),
198                )
199                .await
200                .map_err(|_| PgSearchError::Timeout(TimeoutStage::Enrichment))??;
201                if page.data.len() != expected_rows {
202                    return Err(crate::ExtensionError::rejected(
203                        "typed_page_enricher",
204                        "enrichment must preserve row count",
205                    )
206                    .into());
207                }
208                Ok(page)
209            },
210            |page| page.data.len(),
211        )
212        .await
213    }
214
215    pub async fn search<E>(
216        &self,
217        query: impl QueryInput<E>,
218        scope: &ServerScope<E>,
219    ) -> Result<Page<E>, PgSearchError>
220    where
221        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
222    {
223        self.search_typed(query, scope, None).await
224    }
225
226    /// Executes a page with an additional server-owned HAVING scope.
227    pub async fn search_with_group_scope<E>(
228        &self,
229        query: impl QueryInput<E>,
230        scope: &ServerScope<E>,
231        group_scope: &GroupScope<E>,
232    ) -> Result<Page<E>, PgSearchError>
233    where
234        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
235    {
236        self.search_typed_with_scopes(query, scope, group_scope, None)
237            .await
238    }
239
240    pub async fn search_with<E>(
241        &self,
242        query: impl QueryInput<E>,
243        scope: &ServerScope<E>,
244        options: &SearchOptions<E>,
245    ) -> Result<Page<E>, PgSearchError>
246    where
247        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
248    {
249        self.search_typed(query, scope, Some(options)).await
250    }
251
252    async fn search_typed<E>(
253        &self,
254        query: impl QueryInput<E>,
255        scope: &ServerScope<E>,
256        options: Option<&SearchOptions<E>>,
257    ) -> Result<Page<E>, PgSearchError>
258    where
259        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
260    {
261        self.search_typed_with_scopes(query, scope, &GroupScope::unrestricted(), options)
262            .await
263    }
264
265    /// Executes a typed page with both server scopes and optional summaries.
266    pub async fn search_typed_with_scopes<E>(
267        &self,
268        query: impl QueryInput<E>,
269        scope: &ServerScope<E>,
270        group_scope: &GroupScope<E>,
271        options: Option<&SearchOptions<E>>,
272    ) -> Result<Page<E>, PgSearchError>
273    where
274        E: SearchEntity + Send + Unpin + for<'row> FromRow<'row, PgRow>,
275    {
276        self.observe::<E, _>(
277            "search",
278            async {
279                let plan = self.resolve_plan_with_scopes(
280                    query,
281                    scope,
282                    group_scope,
283                    QueryTarget::TypedPage,
284                    "search",
285                )?;
286                tracing::debug!("query plan ready");
287                if plan.force_empty {
288                    return Ok(Page::empty());
289                }
290                let partition = self.extensions.resolve_partition(&plan, "search")?;
291                let deadline = self.deadline::<E>();
292                if options.is_none()
293                    && self.execution.count_strategy == CountStrategy::WindowWhenPossible
294                    && !matches!(self.execution.strategy, ExecutionStrategy::Snapshot(_))
295                    && self.renderer.supports_window_count(&plan)
296                {
297                    let window = prepare(self.render::<E>(
298                        &plan,
299                        "search",
300                        "window_list",
301                        partition.as_deref(),
302                        RenderRequest::Window,
303                    )?)?;
304                    let fallback = prepare(self.render::<E>(
305                        &plan,
306                        "search",
307                        "window_fallback_count",
308                        partition.as_deref(),
309                        RenderRequest::Cluster(None),
310                    )?)?;
311                    return self
312                        .execute_window_typed::<E, E>(window, fallback, deadline)
313                        .await;
314                }
315                let list = (plan.paging.size > 0)
316                    .then(|| {
317                        self.render::<E>(
318                            &plan,
319                            "search",
320                            "list",
321                            partition.as_deref(),
322                            RenderRequest::List,
323                        )
324                    })
325                    .transpose()?
326                    .map(prepare)
327                    .transpose()?;
328                let cluster = prepare(self.render::<E>(
329                    &plan,
330                    "search",
331                    "cluster",
332                    partition.as_deref(),
333                    RenderRequest::Cluster(options),
334                )?)?;
335                let mut page = self
336                    .execute_typed_pair::<E, E>(list, cluster, options, deadline)
337                    .await?;
338                page.next_cursor = if page.data.len() == plan.paging.size as usize {
339                    self.entity_cursor(&plan, page.data.last())?
340                } else {
341                    None
342                };
343                Ok(page)
344            },
345            |page| page.data.len(),
346        )
347        .await
348    }
349
350    pub async fn search_map<E>(
351        &self,
352        query: impl QueryInput<E>,
353        scope: &ServerScope<E>,
354    ) -> Result<Page<RowMap>, PgSearchError>
355    where
356        E: SearchEntity,
357    {
358        self.search_map_with_options(query, scope, None).await
359    }
360
361    /// Executes a dynamic map page with an additional server-owned HAVING scope.
362    pub async fn search_map_with_group_scope<E>(
363        &self,
364        query: impl QueryInput<E>,
365        scope: &ServerScope<E>,
366        group_scope: &GroupScope<E>,
367    ) -> Result<Page<RowMap>, PgSearchError>
368    where
369        E: SearchEntity,
370    {
371        self.search_map_with_scopes(query, scope, group_scope, None)
372            .await
373    }
374
375    pub async fn search_map_with<E>(
376        &self,
377        query: impl QueryInput<E>,
378        scope: &ServerScope<E>,
379        options: &SearchOptions<E>,
380    ) -> Result<Page<RowMap>, PgSearchError>
381    where
382        E: SearchEntity,
383    {
384        self.search_map_with_options(query, scope, Some(options))
385            .await
386    }
387
388    async fn search_map_with_options<E>(
389        &self,
390        query: impl QueryInput<E>,
391        scope: &ServerScope<E>,
392        options: Option<&SearchOptions<E>>,
393    ) -> Result<Page<RowMap>, PgSearchError>
394    where
395        E: SearchEntity,
396    {
397        self.search_map_with_scopes(query, scope, &GroupScope::unrestricted(), options)
398            .await
399    }
400
401    /// Executes a map page with both server scopes and optional summaries.
402    pub async fn search_map_with_scopes<E>(
403        &self,
404        query: impl QueryInput<E>,
405        scope: &ServerScope<E>,
406        group_scope: &GroupScope<E>,
407        options: Option<&SearchOptions<E>>,
408    ) -> Result<Page<RowMap>, PgSearchError>
409    where
410        E: SearchEntity,
411    {
412        self.observe::<E, _>(
413            "search_map",
414            async {
415                let plan = self.resolve_plan_with_scopes(
416                    query,
417                    scope,
418                    group_scope,
419                    QueryTarget::MapPage,
420                    "search_map",
421                )?;
422                tracing::debug!("query plan ready");
423                if plan.force_empty {
424                    return Ok(Page::empty());
425                }
426                let partition = self.extensions.resolve_partition(&plan, "search_map")?;
427                let deadline = self.deadline::<E>();
428                if options.is_none()
429                    && self.execution.count_strategy == CountStrategy::WindowWhenPossible
430                    && !matches!(self.execution.strategy, ExecutionStrategy::Snapshot(_))
431                    && self.renderer.supports_window_count(&plan)
432                {
433                    let window = prepare(self.render::<E>(
434                        &plan,
435                        "search_map",
436                        "window_list",
437                        partition.as_deref(),
438                        RenderRequest::Window,
439                    )?)?;
440                    let fallback = prepare(self.render::<E>(
441                        &plan,
442                        "search_map",
443                        "window_fallback_count",
444                        partition.as_deref(),
445                        RenderRequest::Cluster(None),
446                    )?)?;
447                    let mut page = self
448                        .execute_window_map(&plan, window, fallback, deadline)
449                        .await?;
450                    self.format_map_rows::<E>(&mut page.data)?;
451                    self.extensions
452                        .filter_rows::<E>("search_map", &mut page.data)?;
453                    return Ok(self
454                        .extensions
455                        .enrich_page(E::META.entity_name, "search_map", page)
456                        .await?);
457                }
458                let list = (plan.paging.size > 0)
459                    .then(|| {
460                        self.render::<E>(
461                            &plan,
462                            "search_map",
463                            "list",
464                            partition.as_deref(),
465                            RenderRequest::List,
466                        )
467                    })
468                    .transpose()?
469                    .map(prepare)
470                    .transpose()?;
471                let cluster = prepare(self.render::<E>(
472                    &plan,
473                    "search_map",
474                    "cluster",
475                    partition.as_deref(),
476                    RenderRequest::Cluster(options),
477                )?)?;
478                let mut page = self
479                    .execute_map_pair(&plan, list, cluster, options, deadline)
480                    .await?;
481                page.next_cursor = if page.data.len() == plan.paging.size as usize {
482                    self.map_cursor(&plan, page.data.last())?
483                } else {
484                    None
485                };
486                self.format_map_rows::<E>(&mut page.data)?;
487                self.extensions
488                    .filter_rows::<E>("search_map", &mut page.data)?;
489                Ok(self
490                    .extensions
491                    .enrich_page(E::META.entity_name, "search_map", page)
492                    .await?)
493            },
494            |page| page.data.len(),
495        )
496        .await
497    }
498}
499
500pub(crate) fn cursor_page_size<E: SearchEntity>(
501    plan: &mut QueryPlan<E>,
502) -> Result<u32, PgSearchError> {
503    if !matches!(plan.paging.mode, sift_query_core::PagingMode::Cursor(_)) {
504        return Err(ParamError::InvalidValue {
505            field: "cursor".to_owned(),
506            reason: "cursor page API requires cursor paging mode".to_owned(),
507        }
508        .into());
509    }
510    let size = plan.paging.size;
511    if size > 0 {
512        plan.paging.size = size.checked_add(1).ok_or_else(|| {
513            ParamError::PagingOutOfRange("cursor page size is too large".to_owned())
514        })?;
515    }
516    Ok(size)
517}