sift_query_sqlx/enrichment/
mod.rs

1use std::fmt;
2pub(crate) mod label;
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::Duration;
7
8use sift_query_core::{EntityMeta, Paging, PredicateArena, QueryPlan, QueryTarget, SearchEntity};
9use thiserror::Error;
10use tokio::time::timeout;
11
12use crate::renderer::hash_shape;
13use crate::{Page, QueryShapeId, RenderedPgQuery, RowMap};
14
15pub type ExtensionFuture<T> = Pin<Box<dyn Future<Output = Result<T, ExtensionError>> + Send>>;
16
17#[derive(Debug, Clone, PartialEq, Eq, Error)]
18#[non_exhaustive]
19pub enum ExtensionError {
20    #[error("invalid PostgreSQL {kind} identifier: {value}")]
21    InvalidIdentifier { kind: &'static str, value: String },
22    #[error("{extension} rejected the query: {message}")]
23    Rejected {
24        extension: &'static str,
25        message: String,
26    },
27    #[error("partition source could not be applied to rendered SQL")]
28    PartitionSourceNotFound,
29    #[error("label loader timed out")]
30    LabelTimeout,
31    #[error("label loader exceeded its id budget of {0}")]
32    LabelIdLimit(usize),
33}
34
35impl ExtensionError {
36    #[must_use]
37    pub fn rejected(extension: &'static str, message: impl Into<String>) -> Self {
38        Self::Rejected {
39            extension,
40            message: message.into(),
41        }
42    }
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct PartitionName {
47    schema: Option<String>,
48    table: String,
49    alias: Option<String>,
50}
51
52impl PartitionName {
53    pub fn new(table: impl Into<String>) -> Result<Self, ExtensionError> {
54        let table = table.into();
55        validate_identifier("table", &table)?;
56        Ok(Self {
57            schema: None,
58            table,
59            alias: None,
60        })
61    }
62
63    pub fn qualified(
64        schema: impl Into<String>,
65        table: impl Into<String>,
66    ) -> Result<Self, ExtensionError> {
67        let schema = schema.into();
68        let table = table.into();
69        validate_identifier("schema", &schema)?;
70        validate_identifier("table", &table)?;
71        Ok(Self {
72            schema: Some(schema),
73            table,
74            alias: None,
75        })
76    }
77
78    pub fn with_alias(mut self, alias: impl Into<String>) -> Result<Self, ExtensionError> {
79        let alias = alias.into();
80        validate_identifier("alias", &alias)?;
81        self.alias = Some(alias);
82        Ok(self)
83    }
84
85    fn render(&self) -> String {
86        let mut output = self.schema.as_ref().map_or_else(String::new, |schema| {
87            format!("{}.", quote_identifier(schema))
88        });
89        output.push_str(&quote_identifier(&self.table));
90        if let Some(alias) = &self.alias {
91            output.push_str(" AS ");
92            output.push_str(&quote_identifier(alias));
93        }
94        output
95    }
96}
97
98pub struct QueryContext<'a> {
99    pub entity: &'static EntityMeta,
100    pub target: QueryTarget,
101    pub operation: &'static str,
102    pub server_predicate: &'a PredicateArena,
103    pub server_group_predicate: &'a PredicateArena,
104    pub row_predicate: &'a PredicateArena,
105    pub group_predicate: &'a PredicateArena,
106    pub paging: &'a Paging,
107}
108
109impl fmt::Debug for QueryContext<'_> {
110    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
111        formatter
112            .debug_struct("QueryContext")
113            .field("entity", &self.entity.entity_name)
114            .field("target", &self.target)
115            .field("operation", &self.operation)
116            .field("page_size", &self.paging.size)
117            .field(
118                "paging_mode",
119                &match self.paging.mode {
120                    sift_query_core::PagingMode::Page { .. } => "page",
121                    sift_query_core::PagingMode::Offset => "offset",
122                    sift_query_core::PagingMode::Cursor(_) => "cursor",
123                },
124            )
125            .finish_non_exhaustive()
126    }
127}
128
129#[derive(Debug, Clone, Copy)]
130pub struct SqlContext {
131    pub entity_name: &'static str,
132    pub target: QueryTarget,
133    pub operation: &'static str,
134    pub phase: &'static str,
135    pub shape_id: QueryShapeId,
136}
137
138#[derive(Debug, Clone, Copy)]
139pub struct ResultContext {
140    pub entity_name: &'static str,
141    pub operation: &'static str,
142    pub target: QueryTarget,
143}
144
145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
146pub enum ResultFilterDecision {
147    Keep,
148    Discard,
149}
150
151pub trait PartitionResolver: Send + Sync {
152    fn resolve(&self, context: &QueryContext<'_>) -> Result<Option<PartitionName>, ExtensionError>;
153}
154
155pub trait SqlInterceptor: Send + Sync {
156    /// Returns replacement SQL, or `None` to keep the current SQL text.
157    /// Bind values are intentionally not exposed to interceptors.
158    fn intercept(&self, context: &SqlContext, sql: &str) -> Result<Option<String>, ExtensionError>;
159}
160
161pub trait ResultFilter: Send + Sync {
162    /// Mutates or discards a decoded dynamic row. Database `total_count` is not
163    /// changed when a row is discarded.
164    fn filter(
165        &self,
166        context: &ResultContext,
167        row_index: usize,
168        row: &mut RowMap,
169    ) -> Result<ResultFilterDecision, ExtensionError>;
170}
171
172/// Context passed to post-query, whole-page enrichers.
173#[derive(Debug, Clone, Copy)]
174pub struct EnrichContext {
175    pub entity_name: &'static str,
176    pub operation: &'static str,
177}
178
179/// Asynchronous page-level result enhancement. Implementations should batch
180/// external lookups and must not change count, summaries, or cursor metadata.
181pub trait PageEnricher: Send + Sync {
182    fn enrich(&self, page: Page<RowMap>, context: EnrichContext) -> ExtensionFuture<Page<RowMap>>;
183}
184
185/// Explicit typed opt-in: implementers declare which presentation fields may
186/// be filled and preserve database fields and ordering keys.
187pub trait TypedPageEnricher<T>: Send + Sync {
188    fn enrich(&self, rows: Vec<T>, context: EnrichContext) -> ExtensionFuture<Vec<T>>;
189}
190
191#[derive(Clone, Default)]
192pub(crate) struct PgExtensions {
193    partition_resolver: Option<Arc<dyn PartitionResolver>>,
194    sql_interceptors: Vec<Arc<dyn SqlInterceptor>>,
195    result_filters: Vec<Arc<dyn ResultFilter>>,
196    page_enrichers: Vec<Arc<dyn PageEnricher>>,
197}
198
199impl fmt::Debug for PgExtensions {
200    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
201        formatter
202            .debug_struct("PgExtensions")
203            .field("has_partition_resolver", &self.partition_resolver.is_some())
204            .field("sql_interceptors", &self.sql_interceptors.len())
205            .field("result_filters", &self.result_filters.len())
206            .field("page_enrichers", &self.page_enrichers.len())
207            .finish()
208    }
209}
210
211impl PgExtensions {
212    pub(crate) fn set_partition_resolver(&mut self, resolver: Arc<dyn PartitionResolver>) {
213        self.partition_resolver = Some(resolver);
214    }
215
216    pub(crate) fn push_sql_interceptor(&mut self, interceptor: Arc<dyn SqlInterceptor>) {
217        self.sql_interceptors.push(interceptor);
218    }
219
220    pub(crate) fn push_result_filter(&mut self, filter: Arc<dyn ResultFilter>) {
221        self.result_filters.push(filter);
222    }
223
224    pub(crate) fn push_page_enricher(&mut self, enricher: Arc<dyn PageEnricher>) {
225        self.page_enrichers.push(enricher);
226    }
227
228    pub(crate) async fn enrich_page(
229        &self,
230        entity_name: &'static str,
231        operation: &'static str,
232        mut page: Page<RowMap>,
233    ) -> Result<Page<RowMap>, ExtensionError> {
234        let context = EnrichContext {
235            entity_name,
236            operation,
237        };
238        for enricher in &self.page_enrichers {
239            let count = page.total_count;
240            let summaries = page.summaries.clone();
241            let cursor = page.next_cursor.clone();
242            let rows = page.data.len();
243            page = timeout(Duration::from_secs(5), enricher.enrich(page, context))
244                .await
245                .map_err(|_| ExtensionError::LabelTimeout)??;
246            if page.total_count != count
247                || page.summaries != summaries
248                || page.next_cursor != cursor
249                || page.data.len() != rows
250            {
251                return Err(ExtensionError::rejected(
252                    "page_enricher",
253                    "enrichment must preserve row count, count, summaries and cursor",
254                ));
255            }
256        }
257        Ok(page)
258    }
259
260    pub(crate) fn resolve_partition<E: SearchEntity>(
261        &self,
262        plan: &QueryPlan<E>,
263        operation: &'static str,
264    ) -> Result<Option<String>, ExtensionError> {
265        let context = QueryContext {
266            entity: plan.meta,
267            target: plan.target,
268            operation,
269            server_predicate: &plan.server_predicate,
270            server_group_predicate: &plan.server_group_predicate,
271            row_predicate: &plan.row_predicate,
272            group_predicate: &plan.group_predicate,
273            paging: &plan.paging,
274        };
275        self.partition_resolver
276            .as_ref()
277            .map(|resolver| resolver.resolve(&context))
278            .transpose()
279            .map(|partition| partition.flatten().map(|partition| partition.render()))
280    }
281
282    pub(crate) fn apply_rendered<E: SearchEntity>(
283        &self,
284        plan: &QueryPlan<E>,
285        operation: &'static str,
286        phase: &'static str,
287        partition: Option<&str>,
288        mut query: RenderedPgQuery,
289    ) -> Result<RenderedPgQuery, ExtensionError> {
290        if partition.is_none() && self.sql_interceptors.is_empty() {
291            return Ok(query);
292        }
293        let mut sql = query.sql.to_string();
294        if let Some(partition) = partition {
295            let source = format!(" FROM {}", plan.meta.from);
296            if !sql.contains(&source) {
297                return Err(ExtensionError::PartitionSourceNotFound);
298            }
299            sql = sql.replace(&source, &format!(" FROM {partition}"));
300        }
301        for interceptor in &self.sql_interceptors {
302            let context = SqlContext {
303                entity_name: plan.meta.entity_name,
304                target: plan.target,
305                operation,
306                phase,
307                shape_id: query.shape_id,
308            };
309            if let Some(replacement) = interceptor.intercept(&context, &sql)? {
310                sql = replacement;
311            }
312        }
313        if sql.as_str() != query.sql.as_ref() {
314            query.shape_id = QueryShapeId(hash_shape(query.shape_id.0, &sql));
315            query.sql = sql.into();
316        }
317        Ok(query)
318    }
319
320    pub(crate) fn filter_rows<E: SearchEntity>(
321        &self,
322        operation: &'static str,
323        rows: &mut Vec<RowMap>,
324    ) -> Result<(), ExtensionError> {
325        if self.result_filters.is_empty() {
326            return Ok(());
327        }
328        let context = ResultContext {
329            entity_name: E::META.entity_name,
330            operation,
331            target: QueryTarget::MapPage,
332        };
333        let mut retained = Vec::with_capacity(rows.len());
334        for (index, mut row) in rows.drain(..).enumerate() {
335            let mut decision = ResultFilterDecision::Keep;
336            for filter in &self.result_filters {
337                if filter.filter(&context, index, &mut row)? == ResultFilterDecision::Discard {
338                    decision = ResultFilterDecision::Discard;
339                    break;
340                }
341            }
342            if decision == ResultFilterDecision::Keep {
343                retained.push(row);
344            }
345        }
346        *rows = retained;
347        Ok(())
348    }
349}
350
351fn validate_identifier(kind: &'static str, value: &str) -> Result<(), ExtensionError> {
352    let mut bytes = value.bytes();
353    let valid = bytes
354        .next()
355        .is_some_and(|first| first.is_ascii_alphabetic() || first == b'_')
356        && bytes.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'$'));
357    if valid {
358        Ok(())
359    } else {
360        Err(ExtensionError::InvalidIdentifier {
361            kind,
362            value: value.to_owned(),
363        })
364    }
365}
366
367fn quote_identifier(value: &str) -> String {
368    format!("\"{value}\"")
369}
370#[cfg(test)]
371mod tests;
sift_query_sqlx/enrichment/mod.rs

sift_query_sqlx/enrichment/
mod.rs
