qql_core/ast/
transform.rs

1use super::{
2    ComparisonOp, FilterExpr, PointId, PointIdPredicate, PointSelector, Prefetch, PrefetchSource,
3    QueryExpr, QueryStmt, ShardKey, Stmt, Value,
4};
5use crate::error::QqlError;
6use alloc::boxed::Box;
7use alloc::string::ToString;
8
9impl Stmt {
10    /// Custom shard routing key for this statement, if any.
11    ///
12    /// Corresponds to QQL `SHARD` on DML, lowered to request-level
13    /// `shard_key` (REST) / `ShardKeySelector` (gRPC) — never inside `Filter`.
14    /// Keyword and numeric forms are preserved (`ShardKey::Number(101)` reads
15    /// back as a number, never coerced to `"101"`).
16    pub fn shard_key(&self) -> Option<&ShardKey> {
17        match self {
18            Self::Query(query) => query.shard_key.as_ref(),
19            Self::Scroll(scroll) => scroll.shard_key.as_ref(),
20            Self::Count(count) => count.shard_key.as_ref(),
21            Self::Facet(facet) => facet.shard_key.as_ref(),
22            Self::Upsert(upsert) => upsert.shard_key.as_ref(),
23            Self::Delete(delete) => delete.shard_key.as_ref(),
24            Self::ClearPayload(clear) => clear.shard_key.as_ref(),
25            Self::DeletePayload(delete) => delete.shard_key.as_ref(),
26            Self::DeleteVector(delete) => delete.shard_key.as_ref(),
27            Self::UpdateVector(update) => update.shard_key.as_ref(),
28            Self::UpdatePayload(update) => update.shard_key.as_ref(),
29            Self::Batch(batch) => {
30                // Members of one batch RPC normally share routing; report a
31                // key only when every member agrees so callers never act on
32                // a partial view.
33                let mut keys = batch.statements.iter().map(Stmt::shard_key);
34                match keys.next() {
35                    None => None,
36                    Some(first) => {
37                        if keys.all(|key| key == first) {
38                            first
39                        } else {
40                            None
41                        }
42                    }
43                }
44            }
45            _ => None,
46        }
47    }
48
49    /// Set custom shard routing (same field as QQL `SHARD`).
50    ///
51    /// Prefer writing the `SHARD` clause in the query when the tenant is known
52    /// at authoring time. Use this setter only when the host resolves the key
53    /// after parse (e.g. from auth context) without re-stringifying QQL.
54    ///
55    /// On `QUERY`, recurses into CTEs and nested prefetch queries so routing
56    /// matches a top-level `SHARD` clause. `None` (or an empty keyword) clears
57    /// the key.
58    /// Returns `false` for statement types that cannot carry routing (DDL, SHOW).
59    pub fn set_shard_key(&mut self, shard_key: Option<ShardKey>) -> bool {
60        let shard_key = shard_key.filter(|k| !matches!(k, ShardKey::Keyword(s) if s.is_empty()));
61        match self {
62            Self::Query(query) => {
63                apply_query_shard(query, shard_key.as_ref());
64                true
65            }
66            Self::Scroll(scroll) => {
67                scroll.shard_key = shard_key;
68                true
69            }
70            Self::Count(count) => {
71                count.shard_key = shard_key;
72                true
73            }
74            Self::Facet(facet) => {
75                facet.shard_key = shard_key;
76                true
77            }
78            Self::Upsert(upsert) => {
79                upsert.shard_key = shard_key;
80                true
81            }
82            Self::Delete(delete) => {
83                delete.shard_key = shard_key;
84                true
85            }
86            Self::ClearPayload(clear) => {
87                clear.shard_key = shard_key;
88                true
89            }
90            Self::DeletePayload(delete) => {
91                delete.shard_key = shard_key;
92                true
93            }
94            Self::DeleteVector(delete) => {
95                delete.shard_key = shard_key;
96                true
97            }
98            Self::UpdateVector(update) => {
99                update.shard_key = shard_key;
100                true
101            }
102            Self::UpdatePayload(update) => {
103                update.shard_key = shard_key;
104                true
105            }
106            Self::Batch(batch) => {
107                // Pre-check before mutating: a mixed batch (e.g. containing
108                // DDL/SHOW) must report failure without half-applying keys to
109                // the members that precede the unsupported one.
110                if !batch.statements.iter().all(can_carry_shard_key) {
111                    return false;
112                }
113                for member in &mut batch.statements {
114                    member.set_shard_key(shard_key.clone());
115                }
116                true
117            }
118            _ => false,
119        }
120    }
121}
122
123/// Whether a statement can carry `SHARD` routing: mirrors the capable arms
124/// of [`Stmt::set_shard_key`]. Used to pre-check batches before mutating so a
125/// mixed batch never half-applies.
126fn can_carry_shard_key(statement: &Stmt) -> bool {
127    match statement {
128        Stmt::Query(_)
129        | Stmt::Scroll(_)
130        | Stmt::Count(_)
131        | Stmt::Facet(_)
132        | Stmt::Upsert(_)
133        | Stmt::Delete(_)
134        | Stmt::ClearPayload(_)
135        | Stmt::DeletePayload(_)
136        | Stmt::DeleteVector(_)
137        | Stmt::UpdateVector(_)
138        | Stmt::UpdatePayload(_) => true,
139        Stmt::Batch(batch) => batch.statements.iter().all(can_carry_shard_key),
140        _ => false,
141    }
142}
143
144/// Apply shard routing to a query and nested CTE / prefetch queries.
145fn apply_query_shard(query: &mut QueryStmt, key: Option<&ShardKey>) {
146    query.shard_key = key.cloned();
147    for cte in &mut query.ctes {
148        apply_query_shard(&mut cte.query, key);
149    }
150    if let Some(prefetches) = expression_prefetch(&mut query.expression) {
151        for prefetch in prefetches {
152            if let PrefetchSource::Query(nested) = &mut prefetch.source {
153                apply_query_shard(nested, key);
154            }
155        }
156    }
157}
158
159/// Injects a typed field comparison into a statement (CTEs and prefetches included), fail-closed.
160///
161/// `QUERY` merges the comparison into the top-level filter, every CTE, and
162/// every prefetch stage; `SCROLL` / `COUNT` / `FACET` merge into their filter;
163/// selector statements (`DELETE`, `CLEAR PAYLOAD`, `DELETE PAYLOAD`,
164/// `DELETE VECTOR`, `UPDATE PAYLOAD`) fold the filter into their point
165/// selector. `BATCH` recurses into every member.
166///
167/// `UPSERT` has no filter to merge into, so the field is stamped onto each
168/// point's payload instead, last-writer-wins: a pre-existing conflicting value
169/// is overwritten, never merged. Overwriting (rather than skipping on
170/// conflict) is what makes tenant stamping sound — a caller-supplied payload
171/// value must not survive the injected policy value. Only `Eq` on a non-`id`
172/// payload field is stampable; anything else fails closed, as does an unbound
173/// whole-point placeholder (`:name` / `?`), which has no payload yet —
174/// silently skipping it would drop a security filter, so callers must bind
175/// point parameters before injection.
176///
177/// `UPDATE VECTOR` is an explicit deny, not an oversight: point-vector
178/// replacement addresses points by ID list and carries no selector, so there
179/// is nothing to merge a tenant filter into. Filter by IDs before building
180/// the statement instead. DDL, `SHOW`, quota, and any other statement type
181/// without a filter or payload surface fail closed with
182/// `QQL-VALIDATION-FILTER-INJECT` rather than silently no-oping, so a policy
183/// bypass can never hide behind an unsupported statement kind.
184pub fn inject_filter(
185    statement: &mut Stmt,
186    field: &str,
187    operator: ComparisonOp,
188    value: Value,
189) -> Result<(), QqlError> {
190    let filter = build_filter(field, operator, value.clone())?;
191    match statement {
192        Stmt::Query(query) => inject_query(query, &filter),
193        Stmt::Batch(batch) => {
194            for member in &mut batch.statements {
195                inject_filter(member, field, operator, value.clone())?;
196            }
197        }
198        Stmt::Scroll(scroll) => merge_filter(&mut scroll.filter, filter),
199        Stmt::Delete(delete) => merge_selector(&mut delete.selector, filter),
200        Stmt::Count(count) => merge_filter(&mut count.filter, filter),
201        Stmt::Facet(facet) => merge_filter(&mut facet.filter, filter),
202        Stmt::ClearPayload(clear) => merge_selector(&mut clear.selector, filter),
203        Stmt::DeletePayload(del) => merge_selector(&mut del.selector, filter),
204        Stmt::DeleteVector(del_vec) => merge_selector(&mut del_vec.selector, filter),
205        Stmt::UpdatePayload(update) => merge_selector(&mut update.selector, filter),
206        // Explicit deny (see rustdoc): no selector exists to merge into.
207        Stmt::UpdateVector(_) => {
208            return Err(QqlError::validation(
209                "QQL-VALIDATION-FILTER-INJECT",
210                "inject_filter does not apply to this statement type (UPDATE VECTOR): point-vector replacement has no selector; address points by ID",
211                None,
212            ));
213        }
214        Stmt::Upsert(_) if operator != ComparisonOp::Eq || field.eq_ignore_ascii_case("id") => {
215            return Err(QqlError::validation(
216                "QQL-VALIDATION-FILTER-INJECT",
217                "inject_filter into UPSERT requires Eq on a non-id payload field",
218                None,
219            ));
220        }
221        Stmt::Upsert(upsert) => {
222            for point in &mut upsert.points {
223                // A whole-point placeholder has no payload yet (see rustdoc).
224                let inline = match point {
225                    crate::ast::PointEntry::Inline(inline) => inline,
226                    crate::ast::PointEntry::Param(name, _) => {
227                        return Err(QqlError::validation(
228                            "QQL-VALIDATION-FILTER-INJECT",
229                            alloc::format!(
230                                "cannot inject filter into unbound point parameter ':{name}'; bind point parameters before filter injection"
231                            ),
232                            None,
233                        ));
234                    }
235                    crate::ast::PointEntry::PositionalParam(idx, _) => {
236                        return Err(QqlError::validation(
237                            "QQL-VALIDATION-FILTER-INJECT",
238                            alloc::format!(
239                                "cannot inject filter into unbound point parameter '?{}'; bind point parameters before filter injection",
240                                *idx + 1
241                            ),
242                            None,
243                        ));
244                    }
245                };
246                if let Some((_, current)) = inline
247                    .payload
248                    .iter_mut()
249                    .find(|(key, _)| key.eq_ignore_ascii_case(field))
250                {
251                    // Last-writer-wins (see rustdoc): overwrite, never merge.
252                    *current = value.clone();
253                } else {
254                    inline.payload.push((field.to_string(), value.clone()));
255                }
256            }
257        }
258        other => {
259            return Err(QqlError::validation(
260                "QQL-VALIDATION-FILTER-INJECT",
261                format!(
262                    "inject_filter does not apply to this statement type ({})",
263                    other.stmt_kind()
264                ),
265                None,
266            ));
267        }
268    }
269    Ok(())
270}
271
272fn build_filter(field: &str, operator: ComparisonOp, value: Value) -> Result<FilterExpr, QqlError> {
273    if field.eq_ignore_ascii_case("id") {
274        if operator != ComparisonOp::Eq {
275            return Err(QqlError::validation(
276                "QQL-VALIDATION-ID-PREDICATE",
277                "point ID injection supports equality only",
278                None,
279            ));
280        }
281        let id = match value {
282            Value::Int(value) if value >= 0 => PointId::Number(value as u64),
283            Value::UInt(value) => PointId::Number(value),
284            Value::Str(value) => PointId::String(value),
285            _ => {
286                return Err(QqlError::validation(
287                    "QQL-VALIDATION-POINT-ID",
288                    "point IDs must be unsigned integers or strings",
289                    None,
290                ));
291            }
292        };
293        Ok(FilterExpr::PointId(PointIdPredicate::Eq(id)))
294    } else {
295        Ok(FilterExpr::Compare {
296            field: field.to_string(),
297            op: operator,
298            value,
299        })
300    }
301}
302
303fn inject_query(query: &mut QueryStmt, filter: &FilterExpr) {
304    merge_filter(&mut query.filter, filter.clone());
305    for cte in &mut query.ctes {
306        inject_query(&mut cte.query, filter);
307    }
308    if let Some(prefetches) = expression_prefetch(&mut query.expression) {
309        for prefetch in prefetches {
310            merge_filter(&mut prefetch.filter, filter.clone());
311            if let PrefetchSource::Query(query) = &mut prefetch.source {
312                inject_query(query, filter);
313            }
314        }
315    }
316}
317
318fn expression_prefetch(expression: &mut QueryExpr) -> Option<&mut Vec<Prefetch>> {
319    match expression {
320        QueryExpr::Nearest { prefetch, .. }
321        | QueryExpr::Recommend { prefetch, .. }
322        | QueryExpr::Context { prefetch, .. }
323        | QueryExpr::Discover { prefetch, .. }
324        | QueryExpr::Fusion { prefetch, .. }
325        | QueryExpr::Formula { prefetch, .. }
326        | QueryExpr::RelevanceFeedback { prefetch, .. }
327        | QueryExpr::Rerank { prefetch, .. }
328        | QueryExpr::CrossRerank { prefetch, .. } => Some(prefetch),
329        QueryExpr::Points { .. }
330        | QueryExpr::OrderBy { .. }
331        | QueryExpr::SampleRandom
332        | QueryExpr::Hybrid { .. } => None,
333    }
334}
335
336fn merge_selector(selector: &mut PointSelector, filter: FilterExpr) {
337    let current = match core::mem::replace(selector, PointSelector::Ids(Vec::new())) {
338        PointSelector::Id(id) => FilterExpr::PointId(PointIdPredicate::Eq(id)),
339        PointSelector::Ids(ids) => FilterExpr::PointId(PointIdPredicate::In(ids)),
340        PointSelector::Filter(existing) => *existing,
341    };
342    *selector = PointSelector::Filter(Box::new(and(current, filter)));
343}
344
345fn merge_filter(current: &mut Option<Box<FilterExpr>>, filter: FilterExpr) {
346    *current = Some(Box::new(match current.take() {
347        Some(current) => and(*current, filter),
348        None => filter,
349    }));
350}
351
352fn and(left: FilterExpr, right: FilterExpr) -> FilterExpr {
353    match left {
354        FilterExpr::And { mut operands } => {
355            operands.push(right);
356            FilterExpr::And { operands }
357        }
358        left => FilterExpr::And {
359            operands: alloc::vec![left, right],
360        },
361    }
362}
qql_core/ast/transform.rs

qql_core/ast/
transform.rs
