nodedb_sql/planner/
const_fold.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Plan-time constant folding for `SqlExpr`.
4//!
5//! Evaluates literal expressions and registered zero-or-few-arg scalar
6//! functions (e.g. `now()`, `current_timestamp`, `date_add(now(), '1h')`)
7//! at plan time via the shared `nodedb_query::functions::eval_function`
8//! evaluator.
9//!
10//! This keeps the bare-`SELECT` projection path, the `INSERT`/`UPSERT`
11//! `VALUES` path, and any future default-expression paths from drifting
12//! apart — they all reach the same evaluator that the Data Plane uses
13//! for column-reference evaluation.
14//!
15//! Semantics: Postgres / SQL-standard compatible. `now()` and
16//! `current_timestamp` snapshot once per statement — `CURRENT_TIMESTAMP`
17//! is defined to return the same value for every row of a single
18//! statement, and Postgres goes further (same value for the whole
19//! transaction). Folding at plan time satisfies both contracts and is
20//! cheaper than per-row runtime dispatch.
21
22use std::sync::LazyLock;
23
24use nodedb_types::Value;
25use sonic_rs;
26
27use crate::functions::registry::{FunctionCategory, FunctionRegistry};
28use crate::types::{BinaryOp, SqlExpr, SqlValue, UnaryOp};
29
30/// Process-wide default registry. Used by call sites that don't already
31/// thread a `FunctionRegistry` through (e.g. the DML `VALUES` path).
32static DEFAULT_REGISTRY: LazyLock<FunctionRegistry> = LazyLock::new(FunctionRegistry::new);
33
34/// Access the shared default registry.
35pub fn default_registry() -> &'static FunctionRegistry {
36    &DEFAULT_REGISTRY
37}
38
39/// Convenience wrapper around [`fold_constant`] using the default registry.
40pub fn fold_constant_default(expr: &SqlExpr) -> Option<SqlValue> {
41    fold_constant(expr, default_registry())
42}
43
44/// Fold a `SqlExpr` to a literal `SqlValue` at plan time, or return
45/// `None` if the expression depends on row/runtime state (column refs,
46/// subqueries, unknown functions, etc.).
47pub fn fold_constant(expr: &SqlExpr, registry: &FunctionRegistry) -> Option<SqlValue> {
48    match expr {
49        SqlExpr::Literal(v) => Some(v.clone()),
50        SqlExpr::ArrayLiteral(items) => items
51            .iter()
52            .map(|item| fold_constant(item, registry))
53            .collect::<Option<Vec<_>>>()
54            .map(SqlValue::Array),
55        SqlExpr::UnaryOp {
56            op: UnaryOp::Neg,
57            expr,
58        } => match fold_constant(expr, registry)? {
59            SqlValue::Int(i) => Some(SqlValue::Int(-i)),
60            SqlValue::Float(f) => Some(SqlValue::Float(-f)),
61            SqlValue::Decimal(d) => Some(SqlValue::Decimal(-d)),
62            _ => None,
63        },
64        SqlExpr::BinaryOp { left, op, right } => {
65            let l = fold_constant(left, registry)?;
66            let r = fold_constant(right, registry)?;
67            fold_binary(l, *op, r)
68        }
69        SqlExpr::Function { name, args, .. } => fold_function_call(name, args, registry),
70        SqlExpr::Cast { expr, to_type } => fold_cast(fold_constant(expr, registry)?, to_type),
71        _ => None,
72    }
73}
74
75/// Fold a CAST at plan time. Only applies when the inner expression is already
76/// a constant. The `to_type` string comes from sqlparser's `format!("{data_type}")`
77/// output, so parameterised types like `NUMERIC(5,1)` must be matched by prefix.
78fn fold_cast(inner: SqlValue, to_type: &str) -> Option<SqlValue> {
79    let upper = to_type.to_uppercase();
80    // Strip any precision/scale suffix: "NUMERIC(5,1)" → "NUMERIC".
81    let base = upper
82        .split('(')
83        .next()
84        .map(str::trim)
85        .unwrap_or(upper.as_str());
86
87    match base {
88        // REGCLASS needs the session catalog. Keep the cast intact for the
89        // contextual catalog-fold pass rather than erasing it to a string.
90        "REGCLASS" => None,
91        "NUMERIC" | "DECIMAL" => match inner {
92            SqlValue::Decimal(d) => Some(SqlValue::Decimal(d)),
93            SqlValue::Int(i) => Some(SqlValue::Decimal(rust_decimal::Decimal::from(i))),
94            SqlValue::Float(f) => rust_decimal::Decimal::try_from(f)
95                .ok()
96                .map(SqlValue::Decimal),
97            SqlValue::String(s) => rust_decimal::Decimal::from_str_exact(&s)
98                .ok()
99                .map(SqlValue::Decimal),
100            _ => None,
101        },
102        "INTEGER" | "INT" | "BIGINT" | "SMALLINT" | "INT2" | "INT4" | "INT8" => match inner {
103            SqlValue::Int(i) => Some(SqlValue::Int(i)),
104            SqlValue::Decimal(d) => {
105                rust_decimal::prelude::ToPrimitive::to_i64(&d).map(SqlValue::Int)
106            }
107            SqlValue::Float(f) => {
108                if f.is_finite() {
109                    Some(SqlValue::Int(f as i64))
110                } else {
111                    None
112                }
113            }
114            SqlValue::String(s) => s.parse::<i64>().ok().map(SqlValue::Int),
115            _ => None,
116        },
117        "FLOAT" | "DOUBLE" | "REAL" | "FLOAT4" | "FLOAT8" | "DOUBLE PRECISION" => match inner {
118            SqlValue::Float(f) => Some(SqlValue::Float(f)),
119            SqlValue::Int(i) => Some(SqlValue::Float(i as f64)),
120            SqlValue::Decimal(d) => {
121                rust_decimal::prelude::ToPrimitive::to_f64(&d).map(SqlValue::Float)
122            }
123            SqlValue::String(s) => s.parse::<f64>().ok().map(SqlValue::Float),
124            _ => None,
125        },
126        "TEXT" | "VARCHAR" | "CHAR" | "CHARACTER VARYING" | "CHARACTER" | "BPCHAR" => match inner {
127            SqlValue::String(s) => Some(SqlValue::String(s)),
128            SqlValue::Int(i) => Some(SqlValue::String(i.to_string())),
129            SqlValue::Float(f) => Some(SqlValue::String(f.to_string())),
130            SqlValue::Decimal(d) => Some(SqlValue::String(d.to_string())),
131            SqlValue::Bool(b) => Some(SqlValue::String(b.to_string())),
132            _ => None,
133        },
134        "BOOL" | "BOOLEAN" => match inner {
135            SqlValue::Bool(b) => Some(SqlValue::Bool(b)),
136            SqlValue::Int(i) => Some(SqlValue::Bool(i != 0)),
137            SqlValue::String(s) => match s.to_lowercase().as_str() {
138                "true" | "t" | "yes" | "1" | "on" => Some(SqlValue::Bool(true)),
139                "false" | "f" | "no" | "0" | "off" => Some(SqlValue::Bool(false)),
140                _ => None,
141            },
142            _ => None,
143        },
144        // `JSON` / `JSONB` — JSON values live internally as their text form
145        // in `SqlValue::String`; the write path parses JSON-looking strings
146        // into document structure. The cast elides to the inner value's JSON
147        // text (mirrors the `::tsvector` / `::tsquery` elision in the resolver).
148        "JSON" | "JSONB" => match inner {
149            SqlValue::String(s) => Some(SqlValue::String(s)),
150            SqlValue::Int(i) => Some(SqlValue::String(i.to_string())),
151            SqlValue::Float(f) => Some(SqlValue::String(f.to_string())),
152            SqlValue::Decimal(d) => Some(SqlValue::String(d.to_string())),
153            SqlValue::Bool(b) => Some(SqlValue::String(b.to_string())),
154            SqlValue::Null => Some(SqlValue::Null),
155            _ => None,
156        },
157        _ => None,
158    }
159}
160
161fn fold_binary(l: SqlValue, op: BinaryOp, r: SqlValue) -> Option<SqlValue> {
162    Some(match (l, op, r) {
163        // Int × Int arithmetic.
164        (SqlValue::Int(a), BinaryOp::Add, SqlValue::Int(b)) => SqlValue::Int(a.checked_add(b)?),
165        (SqlValue::Int(a), BinaryOp::Sub, SqlValue::Int(b)) => SqlValue::Int(a.checked_sub(b)?),
166        (SqlValue::Int(a), BinaryOp::Mul, SqlValue::Int(b)) => SqlValue::Int(a.checked_mul(b)?),
167        // Float × Float arithmetic.
168        (SqlValue::Float(a), BinaryOp::Add, SqlValue::Float(b)) => SqlValue::Float(a + b),
169        (SqlValue::Float(a), BinaryOp::Sub, SqlValue::Float(b)) => SqlValue::Float(a - b),
170        (SqlValue::Float(a), BinaryOp::Mul, SqlValue::Float(b)) => SqlValue::Float(a * b),
171        // Decimal × Decimal arithmetic.
172        (SqlValue::Decimal(a), BinaryOp::Add, SqlValue::Decimal(b)) => {
173            SqlValue::Decimal(a.checked_add(b)?)
174        }
175        (SqlValue::Decimal(a), BinaryOp::Sub, SqlValue::Decimal(b)) => {
176            SqlValue::Decimal(a.checked_sub(b)?)
177        }
178        (SqlValue::Decimal(a), BinaryOp::Mul, SqlValue::Decimal(b)) => {
179            SqlValue::Decimal(a.checked_mul(b)?)
180        }
181        (SqlValue::Decimal(a), BinaryOp::Div, SqlValue::Decimal(b)) => {
182            SqlValue::Decimal(a.checked_div(b)?)
183        }
184        // Decimal × Int widening (Int promotes to Decimal).
185        (SqlValue::Decimal(a), BinaryOp::Add, SqlValue::Int(b)) => {
186            SqlValue::Decimal(a.checked_add(rust_decimal::Decimal::from(b))?)
187        }
188        (SqlValue::Int(a), BinaryOp::Add, SqlValue::Decimal(b)) => {
189            SqlValue::Decimal(rust_decimal::Decimal::from(a).checked_add(b)?)
190        }
191        (SqlValue::Decimal(a), BinaryOp::Sub, SqlValue::Int(b)) => {
192            SqlValue::Decimal(a.checked_sub(rust_decimal::Decimal::from(b))?)
193        }
194        (SqlValue::Int(a), BinaryOp::Sub, SqlValue::Decimal(b)) => {
195            SqlValue::Decimal(rust_decimal::Decimal::from(a).checked_sub(b)?)
196        }
197        (SqlValue::Decimal(a), BinaryOp::Mul, SqlValue::Int(b)) => {
198            SqlValue::Decimal(a.checked_mul(rust_decimal::Decimal::from(b))?)
199        }
200        (SqlValue::Int(a), BinaryOp::Mul, SqlValue::Decimal(b)) => {
201            SqlValue::Decimal(rust_decimal::Decimal::from(a).checked_mul(b)?)
202        }
203        (SqlValue::Decimal(a), BinaryOp::Div, SqlValue::Int(b)) => {
204            SqlValue::Decimal(a.checked_div(rust_decimal::Decimal::from(b))?)
205        }
206        (SqlValue::Int(a), BinaryOp::Div, SqlValue::Decimal(b)) => {
207            SqlValue::Decimal(rust_decimal::Decimal::from(a).checked_div(b)?)
208        }
209        // String concat.
210        (SqlValue::String(a), BinaryOp::Concat, SqlValue::String(b)) => {
211            SqlValue::String(format!("{a}{b}"))
212        }
213        _ => return None,
214    })
215}
216
217/// Fold a function call by recursively folding its arguments, dispatching
218/// through the shared scalar evaluator, and converting the result back to
219/// `SqlValue`. Only folds functions that are present in `registry`, so
220/// callers can distinguish "unknown function" from "known function, all
221/// args folded".
222pub fn fold_function_call(
223    name: &str,
224    args: &[SqlExpr],
225    registry: &FunctionRegistry,
226) -> Option<SqlValue> {
227    // Gate on registry so unknown-function paths keep their existing
228    // fallbacks instead of collapsing to SqlValue::Null. Aggregates and
229    // window functions aren't foldable — they need a row stream.
230    let meta = registry.lookup(name)?;
231    if matches!(
232        meta.category,
233        FunctionCategory::Aggregate | FunctionCategory::Window
234    ) {
235        return None;
236    }
237
238    let folded_args: Vec<Value> = args
239        .iter()
240        .map(|a| fold_constant(a, registry).map(sql_to_ndb_value))
241        .collect::<Option<_>>()?;
242
243    let result = nodedb_query::functions::eval_function(&name.to_lowercase(), &folded_args);
244    Some(ndb_to_sql_value(result))
245}
246
247fn sql_to_ndb_value(v: SqlValue) -> Value {
248    match v {
249        SqlValue::Null => Value::Null,
250        SqlValue::Bool(b) => Value::Bool(b),
251        SqlValue::Int(i) => Value::Integer(i),
252        SqlValue::Float(f) => Value::Float(f),
253        SqlValue::Decimal(d) => Value::Decimal(d),
254        SqlValue::String(s) => Value::String(s),
255        SqlValue::Bytes(b) => Value::Bytes(b),
256        SqlValue::Array(a) => Value::Array(a.into_iter().map(sql_to_ndb_value).collect()),
257        SqlValue::Timestamp(dt) => Value::NaiveDateTime(dt),
258        SqlValue::Timestamptz(dt) => Value::DateTime(dt),
259    }
260}
261
262fn ndb_to_sql_value(v: Value) -> SqlValue {
263    match v {
264        Value::Null => SqlValue::Null,
265        Value::Bool(b) => SqlValue::Bool(b),
266        Value::Integer(i) => SqlValue::Int(i),
267        Value::Float(f) => SqlValue::Float(f),
268        Value::String(s) => SqlValue::String(s),
269        Value::Bytes(b) => SqlValue::Bytes(b),
270        Value::Array(a) => SqlValue::Array(a.into_iter().map(ndb_to_sql_value).collect()),
271        // TZ-aware DateTime → Timestamptz; naive → Timestamp.
272        Value::DateTime(dt) => SqlValue::Timestamptz(dt),
273        Value::NaiveDateTime(dt) => SqlValue::Timestamp(dt),
274        Value::Uuid(s) | Value::Ulid(s) | Value::Regex(s) => SqlValue::String(s),
275        Value::Duration(d) => SqlValue::String(d.to_human()),
276        Value::Decimal(d) => SqlValue::Decimal(d),
277        // Geometry and Object values are serialized to JSON strings so that
278        // nested function calls like ST_Distance(ST_Point(...), ST_Point(...))
279        // survive the SqlValue round-trip. The geo evaluator's geom_arg helper
280        // recovers Geometry from a GeoJSON string; Object results (e.g. from
281        // ST_GeoHashDecode) reach the client as a JSON-encoded string column.
282        Value::Geometry(g) => sonic_rs::to_string(&g)
283            .map(SqlValue::String)
284            .unwrap_or(SqlValue::Null),
285        Value::Object(map) => sonic_rs::to_string(&map)
286            .map(SqlValue::String)
287            .unwrap_or(SqlValue::Null),
288        // Structured and opaque types collapse to Null — callers that
289        // need these go through the runtime expression path, not folding.
290        Value::Set(_) | Value::Range { .. } | Value::Record { .. } | Value::ArrayCell(_) => {
291            SqlValue::Null
292        }
293        // Value is #[non_exhaustive]; future variants collapse to Null in the
294        // constant-folding path — runtime expression evaluation handles them.
295        _ => SqlValue::Null,
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn fold_now_produces_timestamptz() {
305        let registry = FunctionRegistry::new();
306        let expr = SqlExpr::Function {
307            name: "now".into(),
308            args: vec![],
309            distinct: false,
310        };
311        let val = fold_constant(&expr, &registry).expect("now() should fold");
312        match val {
313            SqlValue::Timestamptz(dt) => {
314                // Sanity: must not be epoch (year 1970).
315                assert!(dt.micros > 0, "expected post-epoch timestamp, got micros=0");
316            }
317            other => panic!("expected SqlValue::Timestamptz, got {other:?}"),
318        }
319    }
320
321    #[test]
322    fn fold_current_timestamp_produces_timestamptz() {
323        let registry = FunctionRegistry::new();
324        let expr = SqlExpr::Function {
325            name: "current_timestamp".into(),
326            args: vec![],
327            distinct: false,
328        };
329        assert!(matches!(
330            fold_constant(&expr, &registry),
331            Some(SqlValue::Timestamptz(_))
332        ));
333    }
334
335    #[test]
336    fn fold_unknown_function_returns_none() {
337        let registry = FunctionRegistry::new();
338        let expr = SqlExpr::Function {
339            name: "definitely_not_a_real_function".into(),
340            args: vec![],
341            distinct: false,
342        };
343        assert!(fold_constant(&expr, &registry).is_none());
344    }
345
346    #[test]
347    fn fold_array_literal_recursively() {
348        let registry = FunctionRegistry::new();
349        let expr = SqlExpr::ArrayLiteral(vec![
350            SqlExpr::Literal(SqlValue::String("public".into())),
351            SqlExpr::Literal(SqlValue::Int(42)),
352        ]);
353        assert_eq!(
354            fold_constant(&expr, &registry),
355            Some(SqlValue::Array(vec![
356                SqlValue::String("public".into()),
357                SqlValue::Int(42),
358            ]))
359        );
360    }
361
362    #[test]
363    fn fold_literal_arithmetic_still_works() {
364        let registry = FunctionRegistry::new();
365        let expr = SqlExpr::BinaryOp {
366            left: Box::new(SqlExpr::Literal(SqlValue::Int(2))),
367            op: BinaryOp::Add,
368            right: Box::new(SqlExpr::Literal(SqlValue::Int(3))),
369        };
370        assert_eq!(fold_constant(&expr, &registry), Some(SqlValue::Int(5)));
371    }
372
373    #[test]
374    fn fold_column_ref_returns_none() {
375        let registry = FunctionRegistry::new();
376        let expr = SqlExpr::Column {
377            table: None,
378            name: "name".into(),
379        };
380        assert!(fold_constant(&expr, &registry).is_none());
381    }
382
383    #[test]
384    fn fold_decimal_literal() {
385        let registry = FunctionRegistry::new();
386        let d = rust_decimal::Decimal::new(12345, 2); // 123.45
387        let expr = SqlExpr::Literal(SqlValue::Decimal(d));
388        assert_eq!(fold_constant(&expr, &registry), Some(SqlValue::Decimal(d)));
389    }
390
391    #[test]
392    fn fold_decimal_addition() {
393        use rust_decimal::Decimal;
394        let registry = FunctionRegistry::new();
395        let a = Decimal::new(12345, 2); // 123.45
396        let b = Decimal::new(45678, 2); // 456.78
397        let expr = SqlExpr::BinaryOp {
398            left: Box::new(SqlExpr::Literal(SqlValue::Decimal(a))),
399            op: BinaryOp::Add,
400            right: Box::new(SqlExpr::Literal(SqlValue::Decimal(b))),
401        };
402        let expected = Decimal::new(58023, 2); // 580.23
403        assert_eq!(
404            fold_constant(&expr, &registry),
405            Some(SqlValue::Decimal(expected))
406        );
407    }
408
409    #[test]
410    fn fold_decimal_negation() {
411        use rust_decimal::Decimal;
412        let registry = FunctionRegistry::new();
413        let d = Decimal::new(100, 0);
414        let expr = SqlExpr::UnaryOp {
415            op: crate::types::UnaryOp::Neg,
416            expr: Box::new(SqlExpr::Literal(SqlValue::Decimal(d))),
417        };
418        assert_eq!(fold_constant(&expr, &registry), Some(SqlValue::Decimal(-d)));
419    }
420
421    #[test]
422    fn fold_st_geohash() {
423        let registry = FunctionRegistry::new();
424        let expr = SqlExpr::Function {
425            name: "st_geohash".into(),
426            args: vec![
427                SqlExpr::UnaryOp {
428                    op: UnaryOp::Neg,
429                    expr: Box::new(SqlExpr::Literal(SqlValue::Float(122.4))),
430                },
431                SqlExpr::Literal(SqlValue::Float(37.8)),
432                SqlExpr::Literal(SqlValue::Int(6)),
433            ],
434            distinct: false,
435        };
436        let v = fold_constant(&expr, &registry);
437        match v {
438            Some(SqlValue::String(ref s)) if !s.is_empty() => {}
439            other => panic!("expected non-empty SqlValue::String, got {other:?}"),
440        }
441    }
442
443    #[test]
444    fn fold_st_distance_nested_st_point() {
445        let registry = FunctionRegistry::new();
446        let make_point = |lng: f64, lat: f64| SqlExpr::Function {
447            name: "st_point".into(),
448            args: vec![
449                SqlExpr::Literal(SqlValue::Float(lng)),
450                SqlExpr::Literal(SqlValue::Float(lat)),
451            ],
452            distinct: false,
453        };
454        let expr = SqlExpr::Function {
455            name: "st_distance".into(),
456            args: vec![make_point(-122.4, 37.8), make_point(-87.6, 41.8)],
457            distinct: false,
458        };
459        let v = fold_constant(&expr, &registry);
460        match v {
461            Some(SqlValue::Float(d)) => {
462                assert!(d > 0.0, "distance should be positive, got {d}");
463            }
464            other => panic!("expected SqlValue::Float, got {other:?}"),
465        }
466    }
467
468    #[test]
469    fn fold_decimal_int_widening() {
470        use rust_decimal::Decimal;
471        let registry = FunctionRegistry::new();
472        let d = Decimal::new(100, 0); // 100
473        let expr = SqlExpr::BinaryOp {
474            left: Box::new(SqlExpr::Literal(SqlValue::Decimal(d))),
475            op: BinaryOp::Add,
476            right: Box::new(SqlExpr::Literal(SqlValue::Int(50))),
477        };
478        let expected = Decimal::new(150, 0);
479        assert_eq!(
480            fold_constant(&expr, &registry),
481            Some(SqlValue::Decimal(expected))
482        );
483    }
484}
nodedb_sql/planner/const_fold.rs

nodedb_sql/planner/
const_fold.rs
