Skip to main content

nodedb_sql/planner/
collection_predicate.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Plan a predicate against one collection the way a single-collection
4//! `SELECT ... WHERE` plans its `WHERE`.
5//!
6//! An entry point that receives a predicate as a typed tree instead of SQL
7//! text (the native protocol's metadata filter) plans it here. The
8//! predicate then gets the same column check and the same literal coercion
9//! as the `WHERE` a SQL client writes for the same filter, so both entry
10//! points match the same rows.
11
12use nodedb_types::DatabaseId;
13
14use crate::error::{Result, SqlError};
15use crate::planner::predicate_coerce::coerce_predicate_literals;
16use crate::resolver::columns::{ResolvedTable, TableScope};
17use crate::types::{Filter, FilterExpr, SqlCatalog, SqlExpr};
18
19/// Plan `predicate` as the `WHERE` of `SELECT * FROM <collection>`.
20///
21/// A collection the catalog does not hold is `UnknownTable`. A column the
22/// collection does not have is `UnknownColumn`. A literal compared against a
23/// declared `TIMESTAMP` / `TIMESTAMPTZ` column becomes a typed instant.
24pub fn plan_collection_predicate(
25    catalog: &dyn SqlCatalog,
26    collection: &str,
27    mut predicate: SqlExpr,
28) -> Result<Vec<Filter>> {
29    let info = catalog
30        .resolve_relation(DatabaseId::DEFAULT, collection)?
31        .ok_or_else(|| SqlError::UnknownTable {
32            name: collection.to_string(),
33        })?;
34    let scope = TableScope::single(ResolvedTable {
35        name: collection.to_string(),
36        alias: None,
37        info,
38    })?;
39    check_columns(&predicate, &scope)?;
40    coerce_predicate_literals(&mut predicate, &scope)?;
41    Ok(vec![Filter {
42        expr: FilterExpr::Expr(predicate),
43    }])
44}
45
46/// Reject a column reference that names no column of the scope.
47fn check_columns(expr: &SqlExpr, scope: &TableScope) -> Result<()> {
48    match expr {
49        SqlExpr::Column { table, name } => scope.check_name(table.as_deref(), name),
50        SqlExpr::Literal(_) | SqlExpr::Wildcard | SqlExpr::Subquery(_) => Ok(()),
51        SqlExpr::BinaryOp { left, right, .. } => {
52            check_columns(left, scope)?;
53            check_columns(right, scope)
54        }
55        SqlExpr::UnaryOp { expr, .. }
56        | SqlExpr::Cast { expr, .. }
57        | SqlExpr::IsNull { expr, .. } => check_columns(expr, scope),
58        SqlExpr::Function { args, .. } | SqlExpr::ArrayLiteral(args) => {
59            args.iter().try_for_each(|arg| check_columns(arg, scope))
60        }
61        SqlExpr::Case {
62            operand,
63            when_then,
64            else_expr,
65        } => {
66            if let Some(operand) = operand {
67                check_columns(operand, scope)?;
68            }
69            for (when, then) in when_then {
70                check_columns(when, scope)?;
71                check_columns(then, scope)?;
72            }
73            match else_expr {
74                Some(else_expr) => check_columns(else_expr, scope),
75                None => Ok(()),
76            }
77        }
78        SqlExpr::InList { expr, list, .. } => {
79            check_columns(expr, scope)?;
80            list.iter().try_for_each(|item| check_columns(item, scope))
81        }
82        SqlExpr::Between {
83            expr, low, high, ..
84        } => {
85            check_columns(expr, scope)?;
86            check_columns(low, scope)?;
87            check_columns(high, scope)
88        }
89        SqlExpr::Like { expr, pattern, .. } => {
90            check_columns(expr, scope)?;
91            check_columns(pattern, scope)
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::catalog::SqlCatalogError;
100    use crate::types::{BinaryOp, CollectionInfo, ColumnInfo, EngineType, SqlDataType, SqlValue};
101
102    struct OneCollection;
103
104    impl SqlCatalog for OneCollection {
105        fn get_collection(
106            &self,
107            _database_id: DatabaseId,
108            name: &str,
109        ) -> std::result::Result<Option<CollectionInfo>, SqlCatalogError> {
110            if name != "docs" {
111                return Ok(None);
112            }
113            Ok(Some(CollectionInfo {
114                name: "docs".into(),
115                engine: EngineType::DocumentStrict,
116                columns: vec![ColumnInfo {
117                    name: "category".into(),
118                    data_type: SqlDataType::String,
119                    nullable: true,
120                    is_primary_key: false,
121                    default: None,
122                    raw_type: None,
123                    int_width: None,
124                    float_width: None,
125                }],
126                primary_key: None,
127                has_auto_tier: false,
128                indexes: Vec::new(),
129                bitemporal: false,
130                primary: nodedb_types::PrimaryEngine::Document,
131                vector_primary: None,
132                partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed,
133                open_schema: CollectionInfo::open_schema_for(EngineType::DocumentStrict),
134            }))
135        }
136    }
137
138    fn eq(field: &str, value: &str) -> SqlExpr {
139        SqlExpr::BinaryOp {
140            left: Box::new(SqlExpr::Column {
141                table: None,
142                name: field.into(),
143            }),
144            op: BinaryOp::Eq,
145            right: Box::new(SqlExpr::Literal(SqlValue::String(value.into()))),
146        }
147    }
148
149    #[test]
150    fn declared_column_plans_as_where_expression() {
151        let filters = plan_collection_predicate(&OneCollection, "docs", eq("category", "ai"))
152            .expect("declared column plans");
153        assert_eq!(filters.len(), 1);
154        assert!(matches!(filters[0].expr, FilterExpr::Expr(_)));
155    }
156
157    #[test]
158    fn undeclared_column_is_unknown_column() {
159        let err = plan_collection_predicate(&OneCollection, "docs", eq("ghost", "x"))
160            .expect_err("a strict collection rejects an undeclared column");
161        assert!(matches!(err, SqlError::UnknownColumn { .. }), "{err:?}");
162    }
163
164    #[test]
165    fn unknown_collection_is_unknown_table() {
166        let err = plan_collection_predicate(&OneCollection, "nope", eq("category", "ai"))
167            .expect_err("no such collection");
168        assert!(matches!(err, SqlError::UnknownTable { .. }), "{err:?}");
169    }
170}