Skip to main content

nodedb_sql/planner/array_fn/
maint_fn.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Maintenance ARRAY_* functions: bare `SELECT ARRAY_FLUSH(name)` /
4//! `SELECT ARRAY_COMPACT(name)` with no FROM clause.
5
6use sqlparser::ast;
7
8use super::helpers::{collect_args, require_array_name};
9use crate::error::Result;
10use crate::types::{SqlCatalog, SqlPlan};
11
12/// Try to intercept a no-FROM `SELECT array_flush(name)` /
13/// `SELECT array_compact(name)`. The single projection item must be
14/// a bare function call carrying one string-literal argument.
15pub fn try_plan_array_maint_fn(
16    items: &[ast::SelectItem],
17    catalog: &dyn SqlCatalog,
18) -> Result<Option<SqlPlan>> {
19    if items.len() != 1 {
20        return Ok(None);
21    }
22    let func = match &items[0] {
23        ast::SelectItem::UnnamedExpr(ast::Expr::Function(f))
24        | ast::SelectItem::ExprWithAlias {
25            expr: ast::Expr::Function(f),
26            ..
27        } => f,
28        _ => return Ok(None),
29    };
30    let fn_name = crate::parser::normalize::normalize_object_name_checked(&func.name)?;
31    let arg_exprs = match &func.args {
32        ast::FunctionArguments::List(list) => collect_args(&list.args),
33        _ => Vec::new(),
34    };
35    match fn_name.as_str() {
36        "array_flush" => {
37            let name = require_array_name(&arg_exprs, 0, "ARRAY_FLUSH", catalog)?;
38            Ok(Some(SqlPlan::ArrayFlush { name }))
39        }
40        "array_compact" => {
41            let name = require_array_name(&arg_exprs, 0, "ARRAY_COMPACT", catalog)?;
42            Ok(Some(SqlPlan::ArrayCompact { name }))
43        }
44        _ => Ok(None),
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use crate::catalog::{ArrayCatalogView, SqlCatalogError};
51    use crate::error::Result;
52    use crate::functions::registry::FunctionRegistry;
53    use crate::parser::statement::parse_sql;
54    use crate::types::{CollectionInfo, SqlCatalog, SqlPlan};
55    use crate::types_array::{
56        ArrayAttrAst, ArrayAttrType, ArrayDimAst, ArrayDimType, ArrayDomainBound,
57    };
58
59    struct StubCatalog {
60        view: Option<ArrayCatalogView>,
61        right_view: Option<ArrayCatalogView>,
62    }
63    impl SqlCatalog for StubCatalog {
64        fn get_collection(
65            &self,
66            _: nodedb_types::DatabaseId,
67            _name: &str,
68        ) -> std::result::Result<Option<CollectionInfo>, SqlCatalogError> {
69            Ok(None)
70        }
71        fn lookup_array(&self, name: &str) -> Option<ArrayCatalogView> {
72            if name == "g" || name == "left" {
73                self.view.clone()
74            } else if name == "right" {
75                self.right_view.clone()
76            } else {
77                None
78            }
79        }
80    }
81
82    fn view() -> ArrayCatalogView {
83        ArrayCatalogView {
84            name: "g".into(),
85            dims: vec![
86                ArrayDimAst {
87                    name: "chrom".into(),
88                    dtype: ArrayDimType::Int64,
89                    lo: ArrayDomainBound::Int64(1),
90                    hi: ArrayDomainBound::Int64(23),
91                },
92                ArrayDimAst {
93                    name: "pos".into(),
94                    dtype: ArrayDimType::Int64,
95                    lo: ArrayDomainBound::Int64(0),
96                    hi: ArrayDomainBound::Int64(1_000_000),
97                },
98            ],
99            attrs: vec![
100                ArrayAttrAst {
101                    name: "variant".into(),
102                    dtype: ArrayAttrType::String,
103                    nullable: true,
104                },
105                ArrayAttrAst {
106                    name: "qual".into(),
107                    dtype: ArrayAttrType::Float64,
108                    nullable: true,
109                },
110            ],
111            tile_extents: vec![1, 1_000_000],
112        }
113    }
114
115    fn cat() -> StubCatalog {
116        StubCatalog {
117            view: Some(view()),
118            right_view: Some(view()),
119        }
120    }
121
122    fn plan_one(sql: &str) -> Result<SqlPlan> {
123        let stmts = parse_sql(sql)?;
124        let q = match &stmts[0] {
125            sqlparser::ast::Statement::Query(q) => q,
126            _ => panic!("not a query"),
127        };
128        crate::planner::select::plan_query(
129            q,
130            &cat(),
131            &FunctionRegistry::new(),
132            crate::TemporalScope::default(),
133        )
134    }
135
136    #[test]
137    fn flush_happy() {
138        let p = plan_one("SELECT ARRAY_FLUSH('g')").unwrap();
139        assert!(matches!(p, SqlPlan::ArrayFlush { .. }));
140    }
141
142    #[test]
143    fn compact_happy() {
144        let p = plan_one("SELECT ARRAY_COMPACT('g')").unwrap();
145        assert!(matches!(p, SqlPlan::ArrayCompact { .. }));
146    }
147
148    #[test]
149    fn flush_unknown_array_rejected() {
150        assert!(plan_one("SELECT ARRAY_FLUSH('nope')").is_err());
151    }
152}