Skip to main content

nodedb_sql/planner/dml_helpers/
ast_extract.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! AST extraction helpers shared by the DML planners: pulling a bare table
4//! name out of a join tree, and pulling primary-key point-lookup keys out of
5//! a `WHERE` clause.
6
7use sqlparser::ast;
8
9use super::value_convert::expr_to_sql_value;
10use crate::error::{Result, SqlError};
11use crate::parser::normalize::{normalize_ident, normalize_object_name_checked};
12use crate::types::*;
13
14pub(crate) fn extract_table_name_from_table_with_joins(
15    table: &ast::TableWithJoins,
16) -> Result<String> {
17    match &table.relation {
18        ast::TableFactor::Table { name, .. } => Ok(normalize_object_name_checked(name)?),
19        _ => Err(SqlError::Unsupported {
20            detail: "non-table target in DML".into(),
21        }),
22    }
23}
24
25/// Extract point-operation keys from WHERE clause (WHERE pk = literal OR pk IN (...)).
26pub fn extract_point_keys(selection: Option<&ast::Expr>, info: &CollectionInfo) -> Vec<SqlValue> {
27    let pk = match &info.primary_key {
28        Some(pk) => pk.clone(),
29        None => return Vec::new(),
30    };
31
32    let expr = match selection {
33        Some(e) => e,
34        None => return Vec::new(),
35    };
36
37    let mut keys = Vec::new();
38    // The keys stand for the WHERE only when they cover every row it matches.
39    // A disjunct that is no key equality (`id = 'd' OR v < 2`) matches rows no
40    // key names, so the statement takes the predicate path.
41    if collect_pk_equalities(expr, &pk, &mut keys) {
42        keys
43    } else {
44        Vec::new()
45    }
46}
47
48/// Push the keys `expr` names into `keys`. Returns `false` when `expr` matches
49/// a row no pushed key names: a disjunct that is no key equality, or a key
50/// value that is no literal.
51fn collect_pk_equalities(expr: &ast::Expr, pk: &str, keys: &mut Vec<SqlValue>) -> bool {
52    match expr {
53        ast::Expr::BinaryOp {
54            left,
55            op: ast::BinaryOperator::Eq,
56            right,
57        } => {
58            let value = if is_column(left, pk) {
59                right
60            } else if is_column(right, pk) {
61                left
62            } else {
63                return false;
64            };
65            match expr_to_sql_value(value) {
66                Ok(v) => {
67                    keys.push(v);
68                    true
69                }
70                Err(_) => false,
71            }
72        }
73        ast::Expr::BinaryOp {
74            left,
75            op: ast::BinaryOperator::Or,
76            right,
77        } => collect_pk_equalities(left, pk, keys) && collect_pk_equalities(right, pk, keys),
78        ast::Expr::InList {
79            expr: inner,
80            list,
81            negated: false,
82        } if is_column(inner, pk) => list.iter().all(|item| match expr_to_sql_value(item) {
83            Ok(v) => {
84                keys.push(v);
85                true
86            }
87            Err(_) => false,
88        }),
89        _ => false,
90    }
91}
92
93fn is_column(expr: &ast::Expr, name: &str) -> bool {
94    match expr {
95        ast::Expr::Identifier(ident) => normalize_ident(ident) == name,
96        // Three or more parts: schema.table.col — never matches a plain pk name.
97        ast::Expr::CompoundIdentifier(parts) if parts.len() >= 3 => false,
98        ast::Expr::CompoundIdentifier(parts) if parts.len() == 2 => {
99            normalize_ident(&parts[1]) == name
100        }
101        _ => false,
102    }
103}