Skip to main content

sift_query_sqlx/execution/
prepare.rs

1use super::*;
2
3impl PgSiftQuery {
4    pub fn parse<E: SearchEntity>(
5        &self,
6        raw: &RawParams,
7    ) -> Result<sift_query_core::UserQuery<E>, ParamError> {
8        ParamParser::new(&self.policy)?.parse(raw)
9    }
10
11    pub(crate) fn deadline<E: SearchEntity>(&self) -> Instant {
12        let entity_timeout = Duration::from_millis(E::META.limits.timeout_ms);
13        Instant::now() + self.execution.timeout.min(entity_timeout)
14    }
15
16    pub(crate) fn format_map_rows<E: SearchEntity>(
17        &self,
18        rows: &mut [RowMap],
19    ) -> Result<(), PgSearchError> {
20        let mut error = None;
21        self.policy.for_each_codec_binding::<E>(|field, codec| {
22            if error.is_some() {
23                return;
24            }
25            for row in rows.iter_mut() {
26                if let Some(value) = row.get_mut(field.public_name) {
27                    match codec.format_result(field, value.clone()) {
28                        Ok(formatted) => *value = formatted,
29                        Err(reason) => {
30                            error = Some(crate::DecodeError {
31                                column: field.public_name.to_owned(),
32                                reason,
33                            });
34                            return;
35                        }
36                    }
37                }
38            }
39        })?;
40        if let Some(error) = error {
41            return Err(error.into());
42        }
43        Ok(())
44    }
45
46    pub(crate) fn resolve_plan<E: SearchEntity>(
47        &self,
48        query: impl QueryInput<E>,
49        scope: &ServerScope<E>,
50        target: QueryTarget,
51        operation: &'static str,
52    ) -> Result<QueryPlan<E>, PgSearchError> {
53        self.resolve_plan_with_scopes(query, scope, &GroupScope::unrestricted(), target, operation)
54    }
55
56    pub(crate) fn resolve_plan_with_scopes<E: SearchEntity>(
57        &self,
58        query: impl QueryInput<E>,
59        scope: &ServerScope<E>,
60        group_scope: &GroupScope<E>,
61        target: QueryTarget,
62        operation: &'static str,
63    ) -> Result<QueryPlan<E>, PgSearchError> {
64        let parse_started = Instant::now();
65        let query = query.into_user_query(&self.policy)?;
66        let parse_duration = elapsed_micros(parse_started);
67        let plan_started = Instant::now();
68        let plan = QueryPlan::new_with_scopes(&query, scope, group_scope, target)?;
69        tracing::debug!(
70            entity = E::META.entity_name,
71            operation,
72            parse_duration_us = parse_duration,
73            plan_duration_us = elapsed_micros(plan_started),
74            "planned SiftQuery operation"
75        );
76        Ok(plan)
77    }
78
79    pub(crate) fn render<E: SearchEntity>(
80        &self,
81        plan: &QueryPlan<E>,
82        operation: &'static str,
83        phase: &'static str,
84        partition: Option<&str>,
85        request: RenderRequest<'_, E>,
86    ) -> Result<RenderedPgQuery, PgSearchError> {
87        let started = Instant::now();
88        let (base, cache_hit) = match &self.shape_cache {
89            Some(cache) => cache.render(plan, &self.renderer, &request)?,
90            None => (request.render(&self.renderer, plan)?, false),
91        };
92        let output = self
93            .extensions
94            .apply_rendered(plan, operation, phase, partition, base)?;
95        tracing::Span::current().record("query_shape_id", output.shape_id.0);
96        tracing::debug!(
97            entity = E::META.entity_name,
98            render_phase = phase,
99            query_shape_id = output.shape_id.0,
100            render_duration_us = elapsed_micros(started),
101            shape_cache_hit = cache_hit,
102            "rendered SiftQuery SQL"
103        );
104        Ok(output)
105    }
106
107    pub(crate) async fn acquire(
108        &self,
109        deadline: Instant,
110    ) -> Result<sqlx::pool::PoolConnection<Postgres>, PgSearchError> {
111        let started = Instant::now();
112        let connection = timeout_at(deadline, self.pool.acquire())
113            .await
114            .map_err(|_| PgSearchError::Timeout(TimeoutStage::PoolAcquire))?
115            .map_err(PgSearchError::Sql)?;
116        tracing::debug!(
117            pool_wait_duration_us = elapsed_micros(started),
118            "acquired PostgreSQL connection"
119        );
120        Ok(connection)
121    }
122}
123
124fn elapsed_micros(started: Instant) -> u64 {
125    u64::try_from(started.elapsed().as_micros()).unwrap_or(u64::MAX)
126}