sift_query_sqlx/
buckets.rs

1use sift_query_core::{
2    BucketRequest, BucketResult, QueryInput, QueryTarget, SearchEntity, ServerScope,
3};
4use sqlx::Row;
5use tokio::time::timeout_at;
6
7use crate::execution::io::{configure_statement_timeout, finish_statement, prepare};
8use crate::{PgSearchError, PgSiftQuery, TimeoutStage};
9
10impl PgSiftQuery {
11    /// Executes bounded time-series or numeric buckets with the same filters
12    /// and authorization scope as the source query.
13    pub async fn buckets<E: SearchEntity>(
14        &self,
15        query: impl QueryInput<E>,
16        scope: &ServerScope<E>,
17        requests: &[BucketRequest<E>],
18    ) -> Result<Vec<BucketResult>, PgSearchError> {
19        if requests.len() > 8 {
20            return Err(sift_query_core::ParamError::TooComplex("bucket count").into());
21        }
22        if requests.is_empty() {
23            return Ok(Vec::new());
24        }
25        let plan = self.resolve_plan(query, scope, QueryTarget::Count, "buckets")?;
26        let mut output = Vec::with_capacity(requests.len());
27        for request in requests {
28            if request.schema_fingerprint() != E::META.schema_fingerprint {
29                return Err(sift_query_core::ParamError::UnknownField(
30                    "bucket entity schema mismatch".to_owned(),
31                )
32                .into());
33            }
34            let field = E::META.field(request.field_id()).ok_or_else(|| {
35                sift_query_core::ParamError::UnknownField(format!("#{}", request.field_id().0))
36            })?;
37            let rendered = self.renderer.render_bucket(&plan, request)?;
38            let prepared = prepare(rendered)?;
39            let deadline = self.deadline::<E>();
40            let connection = self.acquire(deadline).await?;
41            let mut connection =
42                configure_statement_timeout(connection, deadline, TimeoutStage::ClusterQuery)
43                    .await?;
44            let result = timeout_at(
45                deadline,
46                sqlx::query_with(prepared.sql(), prepared.arguments).fetch_all(&mut *connection),
47            )
48            .await;
49            let Ok(result) = result else {
50                crate::execution::io::close_hard(connection).await;
51                return Err(PgSearchError::Timeout(TimeoutStage::ClusterQuery));
52            };
53            let (mut connection, rows) =
54                finish_statement(connection, result, TimeoutStage::ClusterQuery).await?;
55            let mut buckets = Vec::with_capacity(rows.len());
56            for row in rows {
57                buckets.push(sift_query_core::BucketEntry {
58                    label: row.try_get::<String, _>("sq_bucket")?,
59                    count: row.try_get::<i64, _>("sq_count")?,
60                });
61            }
62            let _ = &mut connection;
63            output.push(BucketResult {
64                field: field.public_name,
65                buckets,
66            });
67        }
68        Ok(output)
69    }
70}
sift_query_sqlx/buckets.rs

sift_query_sqlx/
buckets.rs
