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, ¶ms, &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", ¶m.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, ¶ms, &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}