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#[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#[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 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 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 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 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}