Skip to main content

sift_query_sqlx/
facet.rs

1use sift_query_core::{
2    FacetBucket, FacetRequest, FacetResult, QueryInput, QueryTarget, SearchEntity, ServerScope,
3};
4use sqlx::Row;
5use sqlx::postgres::{PgConnection, PgRow};
6use tokio::time::{Instant, timeout_at};
7
8use crate::execution::io::{configure_statement_timeout, finish_statement, prepare};
9use crate::{PgSearchError, PgSiftQuery, RowMap, TimeoutStage};
10
11/// A page plus server-declared facet buckets.
12#[derive(Debug, Clone, PartialEq)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize))]
14pub struct FacetPage<T> {
15    pub page: crate::Page<T>,
16    pub facets: Vec<FacetResult>,
17}
18
19/// Named result wrapper useful when facets are fetched independently.
20#[derive(Debug, Clone, PartialEq)]
21#[cfg_attr(feature = "serde", derive(serde::Serialize))]
22pub struct FacetResultSet {
23    pub facets: Vec<FacetResult>,
24}
25
26impl PgSiftQuery {
27    /// Executes bounded server-declared facets using the same query and row
28    /// authorization scope. Facet fields and limits are validated before a
29    /// connection is acquired.
30    pub async fn facets<E: SearchEntity>(
31        &self,
32        query: impl QueryInput<E>,
33        scope: &ServerScope<E>,
34        requests: &[FacetRequest],
35    ) -> Result<Vec<FacetResult>, PgSearchError> {
36        if requests.len() > 8 {
37            return Err(sift_query_core::ParamError::TooComplex("facet count").into());
38        }
39        if requests.is_empty() {
40            return Ok(Vec::new());
41        }
42        let plan = self.resolve_plan(query, scope, QueryTarget::Count, "facets")?;
43        let prepared = self.prepare_facet_queries(&plan, requests)?;
44        let mut output = Vec::with_capacity(prepared.len());
45        let deadline = self.deadline::<E>();
46        for (prepared, field, value_type) in prepared {
47            let connection = self.acquire(deadline).await?;
48            output.push(fetch_facet(connection, prepared, field, value_type, deadline).await?);
49        }
50        Ok(output)
51    }
52
53    /// Convenience operation for dynamic list pages with facets.
54    pub async fn search_map_with_facets<E: SearchEntity>(
55        &self,
56        query: impl QueryInput<E> + Clone,
57        scope: &ServerScope<E>,
58        requests: &[FacetRequest],
59    ) -> Result<FacetPage<RowMap>, PgSearchError> {
60        let page = self.search_map(query.clone(), scope).await?;
61        let facets = self.facets(query, scope, requests).await?;
62        Ok(FacetPage { page, facets })
63    }
64
65    /// Executes a facet on a caller-owned connection. This is useful when the
66    /// caller has opened a read-only transaction and needs the list and facet
67    /// queries to observe one snapshot.
68    pub async fn facets_on<E: SearchEntity>(
69        &self,
70        connection: &mut PgConnection,
71        query: impl QueryInput<E>,
72        scope: &ServerScope<E>,
73        requests: &[FacetRequest],
74    ) -> Result<Vec<FacetResult>, PgSearchError> {
75        if requests.is_empty() {
76            return Ok(Vec::new());
77        }
78        self.validate_external_strategy()?;
79        let plan = self.resolve_plan(query, scope, QueryTarget::Count, "facets_on")?;
80        let prepared = self.prepare_facet_queries(&plan, requests)?;
81        let deadline = self.deadline::<E>();
82        let mut output = Vec::with_capacity(prepared.len());
83        for (prepared, field, value_type) in prepared {
84            output.push(
85                fetch_facet_external(connection, prepared, field, value_type, deadline).await?,
86            );
87        }
88        Ok(output)
89    }
90
91    /// Executes a map page and its facets on one caller-owned connection.
92    /// Begin a read-only transaction before calling this method to obtain a
93    /// consistent snapshot for both results.
94    pub async fn search_map_with_facets_on<E: SearchEntity>(
95        &self,
96        connection: &mut PgConnection,
97        query: impl QueryInput<E> + Clone,
98        scope: &ServerScope<E>,
99        requests: &[FacetRequest],
100    ) -> Result<FacetPage<RowMap>, PgSearchError> {
101        let page = self.search_map_on(connection, query.clone(), scope).await?;
102        let facets = self.facets_on(connection, query, scope, requests).await?;
103        Ok(FacetPage { page, facets })
104    }
105
106    fn prepare_facet_queries<E: SearchEntity>(
107        &self,
108        plan: &sift_query_core::QueryPlan<E>,
109        requests: &[FacetRequest],
110    ) -> Result<
111        Vec<(
112            crate::execution::io::PreparedPgQuery,
113            &'static str,
114            sift_query_core::ValueType,
115        )>,
116        PgSearchError,
117    > {
118        if requests.len() > 8 {
119            return Err(sift_query_core::ParamError::TooComplex("facet count").into());
120        }
121        if requests.is_empty() {
122            return Ok(Vec::new());
123        }
124        if plan.meta.is_grouped() {
125            return Err(sift_query_core::ParamError::InvalidValue {
126                field: "facet".to_owned(),
127                reason: "facets require a non-grouped read model".to_owned(),
128            }
129            .into());
130        }
131        let partition = self.extensions.resolve_partition(plan, "facets")?;
132        let mut output = Vec::with_capacity(requests.len());
133        for request in requests {
134            if request.schema_fingerprint() != plan.meta.schema_fingerprint {
135                return Err(sift_query_core::ParamError::UnknownField(
136                    "facet entity schema mismatch".to_owned(),
137                )
138                .into());
139            }
140            if request.limit() == 0 || request.limit() > 1_000 {
141                return Err(sift_query_core::ParamError::InvalidValue {
142                    field: "facet.limit".to_owned(),
143                    reason: "facet limit must be between 1 and 1000".to_owned(),
144                }
145                .into());
146            }
147            let field = plan.meta.field(request.field_id()).ok_or_else(|| {
148                sift_query_core::ParamError::UnknownField(format!("#{}", request.field_id().0))
149            })?;
150            plan.validate_field_access([request.field_id()])?;
151            if !field.selectable
152                || !field.sortable
153                || field.kind == sift_query_core::FieldKind::FilterOnly
154            {
155                return Err(sift_query_core::ParamError::InvalidValue {
156                    field: field.public_name.to_owned(),
157                    reason: "facet field must be selectable and sortable".to_owned(),
158                }
159                .into());
160            }
161            if !field.value_type.supports_cursor() {
162                return Err(sift_query_core::ParamError::InvalidValue {
163                    field: field.public_name.to_owned(),
164                    reason: "facet field type is not orderable".to_owned(),
165                }
166                .into());
167            }
168            let rendered = self.renderer.render_facet_with_options(
169                plan,
170                request.field_id(),
171                request.limit(),
172                request.includes_null(),
173                request.excludes_self_filter(),
174            )?;
175            let rendered = self.extensions.apply_rendered(
176                plan,
177                "facets",
178                "facet",
179                partition.as_deref(),
180                rendered,
181            )?;
182            output.push((prepare(rendered)?, field.public_name, field.value_type));
183        }
184        Ok(output)
185    }
186}
187
188async fn fetch_facet(
189    connection: sqlx::pool::PoolConnection<sqlx::Postgres>,
190    prepared: crate::execution::io::PreparedPgQuery,
191    field: &'static str,
192    value_type: sift_query_core::ValueType,
193    deadline: Instant,
194) -> Result<FacetResult, PgSearchError> {
195    let mut connection =
196        configure_statement_timeout(connection, deadline, TimeoutStage::ClusterQuery).await?;
197    let result = timeout_at(
198        deadline,
199        sqlx::query_with(prepared.sql(), prepared.arguments).fetch_all(&mut *connection),
200    )
201    .await;
202    let Ok(result) = result else {
203        crate::execution::io::close_hard(connection).await;
204        return Err(PgSearchError::Timeout(TimeoutStage::ClusterQuery));
205    };
206    let (_connection, rows) =
207        finish_statement(connection, result, TimeoutStage::ClusterQuery).await?;
208    decode_facet_rows(&rows, field, value_type)
209}
210
211async fn fetch_facet_external(
212    connection: &mut PgConnection,
213    prepared: crate::execution::io::PreparedPgQuery,
214    field: &'static str,
215    value_type: sift_query_core::ValueType,
216    deadline: Instant,
217) -> Result<FacetResult, PgSearchError> {
218    let result = timeout_at(
219        deadline,
220        sqlx::query_with(prepared.sql(), prepared.arguments).fetch_all(&mut *connection),
221    )
222    .await;
223    let rows = result.map_err(|_| PgSearchError::Timeout(TimeoutStage::ClusterQuery))??;
224    decode_facet_rows(&rows, field, value_type)
225}
226
227fn decode_facet_rows(
228    rows: &[PgRow],
229    field: &'static str,
230    value_type: sift_query_core::ValueType,
231) -> Result<FacetResult, PgSearchError> {
232    let mut buckets = Vec::with_capacity(rows.len());
233    let mut represented = 0_i64;
234    let mut total = 0_i64;
235    for row in rows {
236        let value = crate::postgres::decode::decode_row_value(row, "sq_key", value_type)?;
237        let count = row.try_get::<i64, _>("sq_count")?;
238        total = row.try_get::<i64, _>("sq_total")?;
239        represented = represented.saturating_add(count);
240        buckets.push(FacetBucket { value, count });
241    }
242    Ok(FacetResult {
243        field,
244        buckets,
245        other_count: total.saturating_sub(represented).max(0),
246    })
247}