Skip to main content

sift_query_sqlx/execution/
projection.rs

1use sift_query_core::{
2    QueryInput, QueryPlan, SearchEntity, SearchOptions, SearchProjection, ServerScope,
3};
4use sqlx::FromRow;
5use sqlx::postgres::PgRow;
6
7use crate::cache::RenderRequest;
8use crate::execution::io::{fetch_typed, prepare};
9use crate::{CountStrategy, CursorPage, ExecutionStrategy, Page, PgSearchError, PgSiftQuery};
10
11impl PgSiftQuery {
12    /// Executes a static typed projection as a cursor page without a count.
13    pub async fn search_as_cursor<E, P>(
14        &self,
15        query: impl QueryInput<E>,
16        scope: &ServerScope<E>,
17    ) -> Result<CursorPage<P>, PgSearchError>
18    where
19        E: SearchEntity,
20        P: SearchProjection<E> + Send + Unpin + for<'row> FromRow<'row, PgRow>,
21    {
22        self.observe::<E, _>(
23            "search_as_cursor",
24            async {
25                let query = query.into_user_query(&self.policy)?;
26                let mut plan = QueryPlan::for_projection::<P>(&query, scope)?;
27                let size = super::search::cursor_page_size(&mut plan)?;
28                if plan.force_empty || size == 0 {
29                    return Ok(CursorPage::empty());
30                }
31                let partition = self
32                    .extensions
33                    .resolve_partition(&plan, "search_as_cursor")?;
34                let rendered = prepare(self.render::<E>(
35                    &plan,
36                    "search_as_cursor",
37                    "static_projection",
38                    partition.as_deref(),
39                    RenderRequest::List,
40                )?)?;
41                let deadline = self.deadline::<E>();
42                let connection = self.acquire(deadline).await?;
43                let (_, mut data) = fetch_typed::<P>(
44                    connection,
45                    rendered,
46                    deadline,
47                    crate::TimeoutStage::ListQuery,
48                )
49                .await?;
50                let has_more = data.len() > size as usize;
51                if has_more {
52                    data.truncate(size as usize);
53                }
54                let next_cursor = if has_more {
55                    self.projection_cursor::<E, P>(&plan, data.last())?
56                } else {
57                    None
58                };
59                Ok(CursorPage {
60                    data,
61                    next_cursor,
62                    has_more,
63                })
64            },
65            |page| page.data.len(),
66        )
67        .await
68    }
69
70    pub async fn search_as<E, P>(
71        &self,
72        query: impl QueryInput<E>,
73        scope: &ServerScope<E>,
74    ) -> Result<Page<P>, PgSearchError>
75    where
76        E: SearchEntity,
77        P: SearchProjection<E> + Send + Unpin + for<'row> FromRow<'row, PgRow>,
78    {
79        self.search_as_with_options(query, scope, None).await
80    }
81
82    pub async fn search_as_with<E, P>(
83        &self,
84        query: impl QueryInput<E>,
85        scope: &ServerScope<E>,
86        options: &SearchOptions<E>,
87    ) -> Result<Page<P>, PgSearchError>
88    where
89        E: SearchEntity,
90        P: SearchProjection<E> + Send + Unpin + for<'row> FromRow<'row, PgRow>,
91    {
92        self.search_as_with_options(query, scope, Some(options))
93            .await
94    }
95
96    async fn search_as_with_options<E, P>(
97        &self,
98        query: impl QueryInput<E>,
99        scope: &ServerScope<E>,
100        options: Option<&SearchOptions<E>>,
101    ) -> Result<Page<P>, PgSearchError>
102    where
103        E: SearchEntity,
104        P: SearchProjection<E> + Send + Unpin + for<'row> FromRow<'row, PgRow>,
105    {
106        self.observe::<E, _>(
107            "search_as",
108            async {
109                let query = query.into_user_query(&self.policy)?;
110                let plan = QueryPlan::for_projection::<P>(&query, scope)?;
111                if plan.force_empty {
112                    return Ok(Page::empty());
113                }
114                let partition = self.extensions.resolve_partition(&plan, "search_as")?;
115                let deadline = self.deadline::<E>();
116                if options.is_none()
117                    && self.execution.count_strategy == CountStrategy::WindowWhenPossible
118                    && !matches!(self.execution.strategy, ExecutionStrategy::Snapshot(_))
119                    && self.renderer.supports_window_count(&plan)
120                {
121                    let window = prepare(self.render::<E>(
122                        &plan,
123                        "search_as",
124                        "window_list",
125                        partition.as_deref(),
126                        RenderRequest::Window,
127                    )?)?;
128                    let fallback = prepare(self.render::<E>(
129                        &plan,
130                        "search_as",
131                        "window_fallback_count",
132                        partition.as_deref(),
133                        RenderRequest::Cluster(None),
134                    )?)?;
135                    return self
136                        .execute_window_typed::<E, P>(window, fallback, deadline)
137                        .await;
138                }
139                let list = (plan.paging.size > 0)
140                    .then(|| {
141                        self.render::<E>(
142                            &plan,
143                            "search_as",
144                            "static_projection",
145                            partition.as_deref(),
146                            RenderRequest::List,
147                        )
148                    })
149                    .transpose()?
150                    .map(prepare)
151                    .transpose()?;
152                let cluster = prepare(self.render::<E>(
153                    &plan,
154                    "search_as",
155                    "cluster",
156                    partition.as_deref(),
157                    RenderRequest::Cluster(options),
158                )?)?;
159                let mut page = self
160                    .execute_typed_pair::<E, P>(list, cluster, options, deadline)
161                    .await?;
162                page.next_cursor =
163                    if matches!(plan.paging.mode, sift_query_core::PagingMode::Cursor(_))
164                        && page.data.len() == plan.paging.size as usize
165                    {
166                        self.projection_cursor(&plan, page.data.last())?
167                    } else {
168                        None
169                    };
170                Ok(page)
171            },
172            |page| page.data.len(),
173        )
174        .await
175    }
176}