Skip to main content

mongreldb_core/
procedure.rs

1use crate::error::{MongrelError, Result};
2use crate::rowid::RowId;
3use crate::schema::TypeId;
4use crate::Value;
5use serde::{Deserialize, Serialize};
6use sha2::{Digest, Sha256};
7use std::collections::{HashMap, HashSet};
8
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum ProcedureMode {
12    ReadOnly,
13    ReadWrite,
14}
15
16#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
17pub struct ProcedureParam {
18    pub name: String,
19    pub ty: TypeId,
20    pub nullable: bool,
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub default: Option<Value>,
23}
24
25#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
26pub struct ProcedureBody {
27    pub steps: Vec<ProcedureStep>,
28    pub return_value: ProcedureValue,
29}
30
31#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
32#[serde(rename_all = "snake_case", tag = "kind")]
33pub enum ProcedureStep {
34    NativeQuery {
35        id: String,
36        table: String,
37        #[serde(default)]
38        conditions: Vec<ProcedureCondition>,
39        #[serde(default, skip_serializing_if = "Option::is_none")]
40        projection: Option<Vec<u16>>,
41        #[serde(default, skip_serializing_if = "Option::is_none")]
42        limit: Option<usize>,
43    },
44    Put {
45        id: String,
46        table: String,
47        cells: Vec<ProcedureCell>,
48        #[serde(default)]
49        returning: bool,
50    },
51    Upsert {
52        id: String,
53        table: String,
54        cells: Vec<ProcedureCell>,
55        #[serde(default, skip_serializing_if = "Option::is_none")]
56        update_cells: Option<Vec<ProcedureCell>>,
57        #[serde(default)]
58        returning: bool,
59    },
60    DeleteByPk {
61        id: String,
62        table: String,
63        pk: ProcedureValue,
64    },
65    DeleteRows {
66        id: String,
67        table: String,
68        row_ids: ProcedureValue,
69    },
70    SqlQuery {
71        id: String,
72        sql: String,
73        #[serde(default)]
74        params: Vec<ProcedureValue>,
75    },
76}
77
78#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
79pub struct ProcedureCell {
80    pub column_id: u16,
81    pub value: ProcedureValue,
82}
83
84#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case", tag = "kind")]
86pub enum ProcedureCondition {
87    Pk {
88        value: ProcedureValue,
89    },
90    BitmapEq {
91        column_id: u16,
92        value: ProcedureValue,
93    },
94    BitmapIn {
95        column_id: u16,
96        values: Vec<ProcedureValue>,
97    },
98    Range {
99        column_id: u16,
100        lo: ProcedureValue,
101        hi: ProcedureValue,
102    },
103    RangeF64 {
104        column_id: u16,
105        lo: ProcedureValue,
106        lo_inclusive: bool,
107        hi: ProcedureValue,
108        hi_inclusive: bool,
109    },
110    IsNull {
111        column_id: u16,
112    },
113    IsNotNull {
114        column_id: u16,
115    },
116    FmContains {
117        column_id: u16,
118        pattern: ProcedureValue,
119    },
120}
121
122#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
123#[serde(rename_all = "snake_case", tag = "kind", content = "value")]
124pub enum ProcedureValue {
125    Literal(Value),
126    Param(String),
127    StepRows(String),
128    StepRow(String),
129    StepScalar(String),
130    Object(Vec<(String, ProcedureValue)>),
131    Array(Vec<ProcedureValue>),
132}
133
134#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
135pub struct StoredProcedure {
136    pub name: String,
137    pub version: u64,
138    pub mode: ProcedureMode,
139    pub params: Vec<ProcedureParam>,
140    pub body: ProcedureBody,
141    pub checksum: String,
142    pub created_epoch: u64,
143    pub updated_epoch: u64,
144}
145
146#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
147pub struct ProcedureEntry {
148    pub procedure: StoredProcedure,
149}
150
151#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
152pub struct ProcedureCallResult {
153    #[serde(default, skip_serializing_if = "Option::is_none")]
154    pub epoch: Option<u64>,
155    pub output: ProcedureCallOutput,
156}
157
158#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
159#[serde(rename_all = "snake_case", tag = "kind", content = "value")]
160pub enum ProcedureCallOutput {
161    Null,
162    Scalar(Value),
163    Row(ProcedureCallRow),
164    Rows(Vec<ProcedureCallRow>),
165    Object(Vec<(String, ProcedureCallOutput)>),
166    Array(Vec<ProcedureCallOutput>),
167}
168
169#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
170pub struct ProcedureCallRow {
171    #[serde(default, skip_serializing_if = "Option::is_none")]
172    pub row_id: Option<RowId>,
173    pub columns: HashMap<u16, Value>,
174}
175
176impl StoredProcedure {
177    pub fn new(
178        name: impl Into<String>,
179        mode: ProcedureMode,
180        params: Vec<ProcedureParam>,
181        body: ProcedureBody,
182        epoch: u64,
183    ) -> Result<Self> {
184        let name = name.into();
185        let checksum = checksum(&name, 1, &mode, &params, &body)?;
186        let proc = Self {
187            name,
188            version: 1,
189            mode,
190            params,
191            body,
192            checksum,
193            created_epoch: epoch,
194            updated_epoch: epoch,
195        };
196        proc.validate()?;
197        Ok(proc)
198    }
199
200    pub fn replaced(&self, mut next: StoredProcedure, epoch: u64) -> Result<StoredProcedure> {
201        next.name = self.name.clone();
202        next.version = self.version + 1;
203        next.created_epoch = self.created_epoch;
204        next.updated_epoch = epoch;
205        next.checksum = checksum(
206            &next.name,
207            next.version,
208            &next.mode,
209            &next.params,
210            &next.body,
211        )?;
212        next.validate()?;
213        Ok(next)
214    }
215
216    pub fn validate(&self) -> Result<()> {
217        validate_name("procedure", &self.name)?;
218        let mut params = HashSet::new();
219        for param in &self.params {
220            validate_name("procedure parameter", &param.name)?;
221            if !params.insert(param.name.as_str()) {
222                return Err(MongrelError::InvalidArgument(format!(
223                    "duplicate procedure parameter {:?}",
224                    param.name
225                )));
226            }
227            if !param.nullable && matches!(param.default, Some(Value::Null)) {
228                return Err(MongrelError::InvalidArgument(format!(
229                    "non-null procedure parameter {:?} has NULL default",
230                    param.name
231                )));
232            }
233        }
234
235        let mut steps = HashSet::new();
236        for step in &self.body.steps {
237            let id = step.id();
238            validate_name("procedure step", id)?;
239            if !steps.insert(id) {
240                return Err(MongrelError::InvalidArgument(format!(
241                    "duplicate procedure step {id:?}"
242                )));
243            }
244            if self.mode == ProcedureMode::ReadOnly && step.is_write() {
245                return Err(MongrelError::InvalidArgument(format!(
246                    "read-only procedure {:?} contains write step {id:?}",
247                    self.name
248                )));
249            }
250        }
251        validate_value_refs(&self.body.return_value, &params, &steps)
252    }
253}
254
255impl ProcedureStep {
256    pub fn id(&self) -> &str {
257        match self {
258            ProcedureStep::NativeQuery { id, .. }
259            | ProcedureStep::Put { id, .. }
260            | ProcedureStep::Upsert { id, .. }
261            | ProcedureStep::DeleteByPk { id, .. }
262            | ProcedureStep::DeleteRows { id, .. }
263            | ProcedureStep::SqlQuery { id, .. } => id,
264        }
265    }
266
267    pub fn table(&self) -> Option<&str> {
268        match self {
269            ProcedureStep::NativeQuery { table, .. }
270            | ProcedureStep::Put { table, .. }
271            | ProcedureStep::Upsert { table, .. }
272            | ProcedureStep::DeleteByPk { table, .. }
273            | ProcedureStep::DeleteRows { table, .. } => Some(table),
274            ProcedureStep::SqlQuery { .. } => None,
275        }
276    }
277
278    pub fn is_write(&self) -> bool {
279        matches!(
280            self,
281            ProcedureStep::Put { .. }
282                | ProcedureStep::Upsert { .. }
283                | ProcedureStep::DeleteByPk { .. }
284                | ProcedureStep::DeleteRows { .. }
285        )
286    }
287}
288
289impl From<StoredProcedure> for ProcedureEntry {
290    fn from(procedure: StoredProcedure) -> Self {
291        Self { procedure }
292    }
293}
294
295#[derive(Serialize)]
296struct ChecksumInput<'a> {
297    name: &'a str,
298    version: u64,
299    mode: &'a ProcedureMode,
300    params: &'a [ProcedureParam],
301    body: &'a ProcedureBody,
302}
303
304fn checksum(
305    name: &str,
306    version: u64,
307    mode: &ProcedureMode,
308    params: &[ProcedureParam],
309    body: &ProcedureBody,
310) -> Result<String> {
311    let input = ChecksumInput {
312        name,
313        version,
314        mode,
315        params,
316        body,
317    };
318    let bytes = serde_json::to_vec(&input)
319        .map_err(|e| MongrelError::Other(format!("procedure checksum serialize: {e}")))?;
320    Ok(hex(&Sha256::digest(bytes)))
321}
322
323fn hex(bytes: &[u8]) -> String {
324    const HEX: &[u8; 16] = b"0123456789abcdef";
325    let mut out = String::with_capacity(bytes.len() * 2);
326    for b in bytes {
327        out.push(HEX[(b >> 4) as usize] as char);
328        out.push(HEX[(b & 0x0f) as usize] as char);
329    }
330    out
331}
332
333fn validate_name(kind: &str, name: &str) -> Result<()> {
334    if name.is_empty()
335        || !name
336            .chars()
337            .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
338    {
339        return Err(MongrelError::InvalidArgument(format!(
340            "{kind} name {name:?} must be non-empty ASCII alphanumeric, '_' or '-'"
341        )));
342    }
343    Ok(())
344}
345
346fn validate_value_refs(
347    value: &ProcedureValue,
348    params: &HashSet<&str>,
349    steps: &HashSet<&str>,
350) -> Result<()> {
351    match value {
352        ProcedureValue::Literal(_) => Ok(()),
353        ProcedureValue::Param(name) => {
354            if params.contains(name.as_str()) {
355                Ok(())
356            } else {
357                Err(MongrelError::InvalidArgument(format!(
358                    "unknown procedure parameter reference {name:?}"
359                )))
360            }
361        }
362        ProcedureValue::StepRows(id)
363        | ProcedureValue::StepRow(id)
364        | ProcedureValue::StepScalar(id) => {
365            if steps.contains(id.as_str()) {
366                Ok(())
367            } else {
368                Err(MongrelError::InvalidArgument(format!(
369                    "unknown procedure step reference {id:?}"
370                )))
371            }
372        }
373        ProcedureValue::Object(fields) => {
374            for (_, v) in fields {
375                validate_value_refs(v, params, steps)?;
376            }
377            Ok(())
378        }
379        ProcedureValue::Array(values) => {
380            for v in values {
381                validate_value_refs(v, params, steps)?;
382            }
383            Ok(())
384        }
385    }
386}