nodedb_sql/planner/
dml_helpers.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use sqlparser::ast;
4
5use crate::error::{Result, SqlError};
6use crate::parser::normalize::{normalize_ident, normalize_object_name_checked};
7use crate::resolver::expr::convert_value;
8use crate::types::*;
9
10pub(super) fn convert_value_rows(
11    columns: &[String],
12    rows: &[Vec<ast::Expr>],
13) -> Result<Vec<Vec<(String, SqlValue)>>> {
14    rows.iter()
15        .map(|row| {
16            row.iter()
17                .enumerate()
18                .map(|(i, expr)| {
19                    let col = columns.get(i).cloned().unwrap_or_else(|| format!("col{i}"));
20                    let val = expr_to_sql_value(expr)?;
21                    Ok((col, val))
22                })
23                .collect::<Result<Vec<_>>>()
24        })
25        .collect()
26}
27
28pub(super) fn expr_to_sql_value(expr: &ast::Expr) -> Result<SqlValue> {
29    match expr {
30        ast::Expr::Value(v) => convert_value(&v.value),
31        // Array literals lower element-wise into `SqlValue::Array`; there is
32        // no array-literal `SqlValue` the constant folder could produce.
33        ast::Expr::Array(ast::Array { elem, .. }) => {
34            let vals = elem.iter().map(expr_to_sql_value).collect::<Result<_>>()?;
35            Ok(SqlValue::Array(vals))
36        }
37        // `ST_Point(...)` / `ST_GeomFromGeoJSON(...)` synthesise a GeoJSON
38        // string in place rather than resolving as registered scalar
39        // functions, so they keep their bespoke handling.
40        ast::Expr::Function(func) => match super::spatial_ctor::try_spatial_constructor(func) {
41            Some(result) => result,
42            // Non-spatial functions (`now()`, `date_add(...)`, registered
43            // scalars) fold through the shared pipeline below.
44            None => fold_constant_value(expr),
45        },
46        // Everything else — `::TYPE` / `CAST(... AS TYPE)` casts, arithmetic,
47        // string concatenation, parenthesised literals — goes through the
48        // same resolver and constant folder the `SELECT` projection path
49        // uses, so the two surfaces never drift. Only genuinely row- or
50        // runtime-dependent expressions (column refs, subqueries, unknown
51        // functions) fail here.
52        _ => fold_constant_value(expr),
53    }
54}
55
56fn fold_constant_value(expr: &ast::Expr) -> Result<SqlValue> {
57    let sql_expr = crate::resolver::expr::convert_expr(expr)?;
58    super::const_fold::fold_constant_default(&sql_expr).ok_or_else(|| SqlError::Unsupported {
59        detail: format!("value expression: {expr}"),
60    })
61}
62
63pub(super) fn extract_table_name_from_table_with_joins(
64    table: &ast::TableWithJoins,
65) -> Result<String> {
66    match &table.relation {
67        ast::TableFactor::Table { name, .. } => Ok(normalize_object_name_checked(name)?),
68        _ => Err(SqlError::Unsupported {
69            detail: "non-table target in DML".into(),
70        }),
71    }
72}
73
74/// Extract point-operation keys from WHERE clause (WHERE pk = literal OR pk IN (...)).
75pub fn extract_point_keys(selection: Option<&ast::Expr>, info: &CollectionInfo) -> Vec<SqlValue> {
76    let pk = match &info.primary_key {
77        Some(pk) => pk.clone(),
78        None => return Vec::new(),
79    };
80
81    let expr = match selection {
82        Some(e) => e,
83        None => return Vec::new(),
84    };
85
86    let mut keys = Vec::new();
87    collect_pk_equalities(expr, &pk, &mut keys);
88    keys
89}
90
91fn collect_pk_equalities(expr: &ast::Expr, pk: &str, keys: &mut Vec<SqlValue>) {
92    match expr {
93        ast::Expr::BinaryOp {
94            left,
95            op: ast::BinaryOperator::Eq,
96            right,
97        } => {
98            if is_column(left, pk)
99                && let Ok(v) = expr_to_sql_value(right)
100            {
101                keys.push(v);
102            } else if is_column(right, pk)
103                && let Ok(v) = expr_to_sql_value(left)
104            {
105                keys.push(v);
106            }
107        }
108        ast::Expr::BinaryOp {
109            left,
110            op: ast::BinaryOperator::Or,
111            right,
112        } => {
113            collect_pk_equalities(left, pk, keys);
114            collect_pk_equalities(right, pk, keys);
115        }
116        ast::Expr::InList {
117            expr: inner,
118            list,
119            negated: false,
120        } if is_column(inner, pk) => {
121            for item in list {
122                if let Ok(v) = expr_to_sql_value(item) {
123                    keys.push(v);
124                }
125            }
126        }
127        _ => {}
128    }
129}
130
131fn is_column(expr: &ast::Expr, name: &str) -> bool {
132    match expr {
133        ast::Expr::Identifier(ident) => normalize_ident(ident) == name,
134        // Three or more parts: schema.table.col — never matches a plain pk name.
135        ast::Expr::CompoundIdentifier(parts) if parts.len() >= 3 => false,
136        ast::Expr::CompoundIdentifier(parts) if parts.len() == 2 => {
137            normalize_ident(&parts[1]) == name
138        }
139        _ => false,
140    }
141}
142
143/// Build a `SqlPlan::VectorPrimaryInsert` from parsed rows.
144///
145/// Extracts the vector-field column into `vector: Vec<f32>` and collects
146/// all remaining columns into `payload_fields`. Rows missing the vector
147/// column are rejected.
148pub(super) fn build_vector_primary_insert_plan(
149    collection: &str,
150    vpc: &nodedb_types::VectorPrimaryConfig,
151    _columns: &[String],
152    rows: Vec<Vec<(String, SqlValue)>>,
153) -> Result<Vec<SqlPlan>> {
154    let mut result_rows = Vec::with_capacity(rows.len());
155    for row in rows {
156        let mut vector: Option<Vec<f32>> = None;
157        let mut payload_fields = std::collections::HashMap::new();
158
159        for (col, val) in row {
160            if col == vpc.vector_field {
161                match val {
162                    SqlValue::Array(items) => {
163                        let floats: Result<Vec<f32>> = items
164                            .iter()
165                            .map(|v| match v {
166                                SqlValue::Float(f) => Ok(*f as f32),
167                                SqlValue::Int(i) => Ok(*i as f32),
168                                SqlValue::Decimal(d) => {
169                                    use rust_decimal::prelude::ToPrimitive;
170                                    d.to_f32().ok_or_else(|| SqlError::Parse {
171                                        detail: format!(
172                                            "vector element decimal '{d}' is out of f32 range"
173                                        ),
174                                    })
175                                }
176                                other => Err(SqlError::Parse {
177                                    detail: format!(
178                                        "vector field must contain numbers, got {other:?}"
179                                    ),
180                                }),
181                            })
182                            .collect();
183                        vector = Some(floats?);
184                    }
185                    other => {
186                        return Err(SqlError::Parse {
187                            detail: format!(
188                                "vector field '{}' must be an array literal, got {other:?}",
189                                vpc.vector_field
190                            ),
191                        });
192                    }
193                }
194            } else {
195                payload_fields.insert(col, val);
196            }
197        }
198
199        let vector = vector.ok_or_else(|| SqlError::Parse {
200            detail: format!(
201                "vector-primary INSERT missing required vector field '{}'",
202                vpc.vector_field
203            ),
204        })?;
205
206        result_rows.push(VectorPrimaryRow {
207            surrogate: nodedb_types::Surrogate::ZERO,
208            vector,
209            payload_fields,
210        });
211    }
212
213    Ok(vec![SqlPlan::VectorPrimaryInsert {
214        collection: collection.to_string(),
215        field: vpc.vector_field.clone(),
216        quantization: vpc.quantization,
217        storage_dtype: vpc.storage_dtype,
218        payload_indexes: vpc.payload_indexes.clone(),
219        rows: result_rows,
220    }])
221}
222
223/// Build a `SqlPlan::KvInsert` from a VALUES clause. Shared by plain INSERT,
224/// UPSERT, and `INSERT ... ON CONFLICT (key) DO UPDATE` — the three paths
225/// differ only in `intent` and `on_conflict_updates`, never in how entries
226/// are extracted from the row exprs.
227///
228/// `pk_col` is the schema-defined primary-key column name from
229/// `CollectionInfo::primary_key`.  When supplied, that column is used as
230/// the KV key regardless of whether it is named `"key"`.  Falls back to
231/// the literal name `"key"` when `pk_col` is `None` (legacy / generic
232/// KV collections that use the built-in key/value column convention).
233pub(super) fn build_kv_insert_plan(
234    table_name: String,
235    columns: &[String],
236    rows_ast: &[Vec<ast::Expr>],
237    intent: KvInsertIntent,
238    on_conflict_updates: Vec<(String, SqlExpr)>,
239    pk_col: Option<&str>,
240) -> Result<Vec<SqlPlan>> {
241    let key_col_name = pk_col.unwrap_or("key");
242    let key_idx = columns.iter().position(|c| c == key_col_name);
243    let ttl_idx = columns.iter().position(|c| c == "ttl");
244    // When using a named primary-key column (e.g. `k STRING PRIMARY KEY`), we
245    // store the key bytes in the KV key slot AND also keep the column in the
246    // value map.  This allows scan filters on the primary-key column (e.g.
247    // `WHERE k = 'x'`) and projection (e.g. `SELECT k FROM ...`) to work
248    // without teaching the KV scan handler to inspect the raw key bytes.
249    // The only column we exclude from the value map is the built-in `"key"`
250    // sentinel (used by raw key/value KV collections) and `"ttl"`.
251    let exclude_from_value: std::collections::HashSet<usize> = {
252        let mut s = std::collections::HashSet::new();
253        // Exclude the raw "key" sentinel column (not a named PK column).
254        if key_col_name == "key"
255            && let Some(idx) = key_idx
256        {
257            s.insert(idx);
258        }
259        if let Some(idx) = ttl_idx {
260            s.insert(idx);
261        }
262        s
263    };
264    let mut entries = Vec::with_capacity(rows_ast.len());
265    let mut ttl_secs: u64 = 0;
266    for row_exprs in rows_ast {
267        let key_val = match key_idx {
268            Some(idx) => expr_to_sql_value(&row_exprs[idx])?,
269            None => SqlValue::String(String::new()),
270        };
271        if let Some(idx) = ttl_idx {
272            match expr_to_sql_value(&row_exprs[idx]) {
273                Ok(SqlValue::Int(n)) => ttl_secs = n.max(0) as u64,
274                Ok(SqlValue::Float(f)) => ttl_secs = f.max(0.0) as u64,
275                _ => {}
276            }
277        }
278        let value_cols: Vec<(String, SqlValue)> = columns
279            .iter()
280            .enumerate()
281            .filter(|(i, _)| !exclude_from_value.contains(i))
282            .map(|(i, col)| {
283                let val = expr_to_sql_value(&row_exprs[i])?;
284                Ok((col.clone(), val))
285            })
286            .collect::<Result<Vec<_>>>()?;
287        entries.push((key_val, value_cols));
288    }
289    Ok(vec![SqlPlan::KvInsert {
290        collection: table_name,
291        entries,
292        ttl_secs,
293        intent,
294        on_conflict_updates,
295    }])
296}
nodedb_sql/planner/dml_helpers.rs

nodedb_sql/planner/
dml_helpers.rs
