nodedb_sql/planner/array_fn/
maint_fn.rs1use sqlparser::ast;
7
8use super::helpers::{collect_args, require_array_name};
9use crate::error::Result;
10use crate::types::{SqlCatalog, SqlPlan};
11
12pub 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}