nodedb_sql/planner/dml_helpers/
kv_counter.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! The row a KV counter atomic (`KV_INCR`, `KV_INCR_FLOAT`) creates for an
4//! absent key.
5//!
6//! The KV engine stores the bytes it is handed and does not know the declared
7//! columns. So the fresh row is planned here, from the catalog, by the same
8//! code a `VALUES` insert runs: `INSERT (key, column) VALUES (key, 0)`. The
9//! Control Plane encodes the cells into the template the engine fills in.
10
11use sqlparser::ast::{self, Expr, Value, ValueWithSpan};
12use sqlparser::tokenizer::Span;
13
14use super::kv_insert::build_kv_insert_plan;
15use super::params::KvInsertParams;
16use crate::catalog::SqlCatalog;
17use crate::error::{Result, SqlError};
18use crate::types::*;
19
20/// The kind of number a counter atomic moves.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum KvCounterKind {
23    /// `KV_INCR` / `KV_DECR`: an integer column.
24    Integer,
25    /// `KV_INCR_FLOAT`: an integer or float column, the same columns it moves
26    /// in an existing row.
27    Float,
28}
29
30/// The row a counter atomic creates for an absent key.
31#[derive(Debug, Clone, PartialEq)]
32pub enum KvCounterFreshRow {
33    /// The collection holds a single `value` column, or declares none: the
34    /// new value is stored as raw decimal text.
35    Raw,
36    /// The collection declares typed columns.
37    Typed {
38        /// The declared column the counter moves: the first column of `kind`
39        /// in name order, never the primary key. An existing row picks its
40        /// column by the same rule. `None` when there is none.
41        column: Option<String>,
42        /// The cells the insert stores other than `column`: DEFAULTs, and the
43        /// primary key when it is a named column.
44        cells: Vec<(String, SqlValue)>,
45    },
46}
47
48/// Plan the row `INSERT (key, column) VALUES (key, 0)` stores in the KV
49/// collection `info`, for a counter of `kind`.
50pub fn plan_kv_counter_fresh_row(
51    info: &CollectionInfo,
52    key: &str,
53    kind: KvCounterKind,
54    catalog: &dyn SqlCatalog,
55) -> Result<KvCounterFreshRow> {
56    if info.engine != EngineType::KeyValue {
57        return Err(SqlError::Unsupported {
58            detail: format!(
59                "KV counter on '{}', which is not a key-value collection",
60                info.name
61            ),
62        });
63    }
64    let pk_col = info.primary_key.as_deref().unwrap_or("key");
65    let value_columns: Vec<&ColumnInfo> = info
66        .columns
67        .iter()
68        .filter(|c| c.name != pk_col && c.name != "key" && c.name != "ttl")
69        .collect();
70    if value_columns.is_empty() || (value_columns.len() == 1 && value_columns[0].name == "value") {
71        return Ok(KvCounterFreshRow::Raw);
72    }
73
74    let column = value_columns
75        .iter()
76        .filter(|c| moves_as(kind, &c.data_type))
77        .map(|c| c.name.clone())
78        .min();
79    let Some(column) = column else {
80        return Ok(KvCounterFreshRow::Typed {
81            column: None,
82            cells: Vec::new(),
83        });
84    };
85
86    let placeholder = match kind {
87        KvCounterKind::Integer => "0",
88        KvCounterKind::Float => "0.0",
89    };
90    let columns = [pk_col.to_string(), column.clone()];
91    let row = ast::Parens::with_empty_span(vec![
92        literal(Value::SingleQuotedString(key.to_string())),
93        literal(Value::Number(placeholder.to_string(), false)),
94    ]);
95    let plans = build_kv_insert_plan(KvInsertParams {
96        collection: info.name.clone(),
97        columns: &columns,
98        rows_ast: std::slice::from_ref(&row),
99        intent: KvInsertIntent::Put,
100        on_conflict_updates: Vec::new(),
101        pk_col: info.primary_key.as_deref(),
102        declared_columns: &info.columns,
103        catalog,
104    })?;
105    let cells = match plans.into_iter().next() {
106        Some(SqlPlan::KvInsert(KvInsertPlan { mut entries, .. })) if entries.len() == 1 => {
107            let (_, cells) = entries.remove(0);
108            cells
109                .into_iter()
110                .filter(|(name, _)| *name != column)
111                .collect()
112        }
113        _ => {
114            return Err(SqlError::Unsupported {
115                detail: "KV counter fresh row did not plan as one KV insert".into(),
116            });
117        }
118    };
119    Ok(KvCounterFreshRow::Typed {
120        column: Some(column),
121        cells,
122    })
123}
124
125/// Whether a declared column of `data_type` is one a counter of `kind`
126/// moves.
127fn moves_as(kind: KvCounterKind, data_type: &SqlDataType) -> bool {
128    match kind {
129        KvCounterKind::Integer => matches!(data_type, SqlDataType::Int64),
130        KvCounterKind::Float => matches!(data_type, SqlDataType::Int64 | SqlDataType::Float64),
131    }
132}
133
134fn literal(value: Value) -> Expr {
135    Expr::Value(ValueWithSpan {
136        value,
137        span: Span::empty(),
138    })
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    /// A catalog with no collections and no sequence state. The DEFAULTs
146    /// these cases declare are literals, so no accessor is reached.
147    struct NoCatalog;
148
149    impl SqlCatalog for NoCatalog {
150        fn get_collection(
151            &self,
152            _database_id: nodedb_types::DatabaseId,
153            _name: &str,
154        ) -> std::result::Result<Option<CollectionInfo>, crate::catalog::SqlCatalogError> {
155            Ok(None)
156        }
157    }
158
159    fn column(name: &str, data_type: SqlDataType, default: Option<&str>) -> ColumnInfo {
160        ColumnInfo {
161            name: name.to_string(),
162            data_type,
163            nullable: true,
164            is_primary_key: false,
165            default: default.map(str::to_string),
166            raw_type: None,
167            int_width: None,
168            float_width: None,
169        }
170    }
171
172    fn kv_info(pk: &str, columns: Vec<ColumnInfo>) -> CollectionInfo {
173        let mut key = column(pk, SqlDataType::String, None);
174        key.is_primary_key = true;
175        let mut all = vec![key];
176        all.extend(columns);
177        CollectionInfo {
178            name: "c".into(),
179            engine: EngineType::KeyValue,
180            columns: all,
181            primary_key: Some(pk.into()),
182            has_auto_tier: false,
183            indexes: Vec::new(),
184            bitemporal: false,
185            primary: nodedb_types::PrimaryEngine::Document,
186            vector_primary: None,
187            partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed,
188            open_schema: CollectionInfo::open_schema_for(EngineType::KeyValue),
189        }
190    }
191
192    #[test]
193    fn a_single_value_column_is_raw() {
194        let info = kv_info("key", vec![column("value", SqlDataType::String, None)]);
195        let row = plan_kv_counter_fresh_row(&info, "k", KvCounterKind::Integer, &NoCatalog)
196            .expect("plan");
197        assert_eq!(row, KvCounterFreshRow::Raw);
198    }
199
200    #[test]
201    fn a_typed_collection_plans_the_insert_row_with_defaults() {
202        let info = kv_info(
203            "key",
204            vec![
205                column("n", SqlDataType::Int64, None),
206                column("status", SqlDataType::String, Some("'new'")),
207                column("score", SqlDataType::Float64, None),
208            ],
209        );
210        let row = plan_kv_counter_fresh_row(&info, "k", KvCounterKind::Integer, &NoCatalog)
211            .expect("plan");
212        let KvCounterFreshRow::Typed { column, cells } = row else {
213            panic!("expected a typed row");
214        };
215        assert_eq!(column.as_deref(), Some("n"));
216        assert!(
217            cells
218                .iter()
219                .any(|(name, value)| name == "status" && *value == SqlValue::String("new".into())),
220            "{cells:?}"
221        );
222        assert!(
223            cells.iter().all(|(name, _)| name != "n" && name != "key"),
224            "{cells:?}"
225        );
226    }
227
228    #[test]
229    fn a_named_primary_key_is_kept_in_the_row() {
230        let info = kv_info("id", vec![column("n", SqlDataType::Int64, None)]);
231        let row = plan_kv_counter_fresh_row(&info, "k1", KvCounterKind::Integer, &NoCatalog)
232            .expect("plan");
233        let KvCounterFreshRow::Typed { cells, .. } = row else {
234            panic!("expected a typed row");
235        };
236        assert!(
237            cells
238                .iter()
239                .any(|(name, value)| name == "id" && *value == SqlValue::String("k1".into())),
240            "{cells:?}"
241        );
242    }
243
244    #[test]
245    fn counters_move_the_first_column_of_their_kind_in_name_order() {
246        let info = kv_info(
247            "key",
248            vec![
249                column("b", SqlDataType::Int64, None),
250                column("a", SqlDataType::Float64, None),
251                column("label", SqlDataType::String, None),
252            ],
253        );
254        let target = |kind| match plan_kv_counter_fresh_row(&info, "k", kind, &NoCatalog) {
255            Ok(KvCounterFreshRow::Typed { column, .. }) => column,
256            other => panic!("expected a typed row, got {other:?}"),
257        };
258        assert_eq!(target(KvCounterKind::Integer).as_deref(), Some("b"));
259        assert_eq!(target(KvCounterKind::Float).as_deref(), Some("a"));
260
261        let text_only = kv_info("key", vec![column("label", SqlDataType::String, None)]);
262        let row = plan_kv_counter_fresh_row(&text_only, "k", KvCounterKind::Integer, &NoCatalog)
263            .expect("plan");
264        assert_eq!(
265            row,
266            KvCounterFreshRow::Typed {
267                column: None,
268                cells: Vec::new(),
269            }
270        );
271    }
272}
nodedb_sql/planner/dml_helpers/kv_counter.rs

nodedb_sql/planner/dml_helpers/
kv_counter.rs
