surrealdb_core/sql/statements/
set.rs

1use crate::ctx::Context;
2use crate::dbs::Options;
3use crate::doc::CursorDoc;
4use crate::err::Error;
5use crate::sql::Value;
6use crate::{cnf::PROTECTED_PARAM_NAMES, sql::Kind};
7
8use reblessive::tree::Stk;
9use revision::revisioned;
10use serde::{Deserialize, Serialize};
11use std::fmt;
12
13#[revisioned(revision = 2)]
14#[derive(Clone, Debug, Default, Eq, PartialEq, PartialOrd, Serialize, Deserialize, Hash)]
15#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
16#[non_exhaustive]
17pub struct SetStatement {
18	pub name: String,
19	pub what: Value,
20	#[revision(start = 2)]
21	pub kind: Option<Kind>,
22}
23
24impl SetStatement {
25	/// Check if we require a writeable transaction
26	pub(crate) fn writeable(&self) -> bool {
27		self.what.writeable()
28	}
29	/// Process this type returning a computed simple Value
30	pub(crate) async fn compute(
31		&self,
32		stk: &mut Stk,
33		ctx: &Context,
34		opt: &Options,
35		doc: Option<&CursorDoc>,
36	) -> Result<Value, Error> {
37		// Check if the variable is a protected variable
38		match PROTECTED_PARAM_NAMES.contains(&self.name.as_str()) {
39			// The variable isn't protected and can be stored
40			false => {
41				let result = self.what.compute(stk, ctx, opt, doc).await?;
42				match self.kind {
43					Some(ref kind) => result
44						.coerce_to(kind)
45						.map_err(|e| e.set_check_from_coerce(self.name.to_string())),
46					None => Ok(result),
47				}
48			}
49			// The user tried to set a protected variable
50			true => Err(Error::InvalidParam {
51				name: self.name.to_owned(),
52			}),
53		}
54	}
55}
56
57impl fmt::Display for SetStatement {
58	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
59		write!(f, "LET ${}", self.name)?;
60		if let Some(ref kind) = self.kind {
61			write!(f, ": {}", kind)?;
62		}
63		write!(f, " = {}", self.what)?;
64		Ok(())
65	}
66}
67
68#[cfg(test)]
69mod tests {
70	use crate::syn::parse;
71
72	#[test]
73	fn check_type() {
74		let query = parse("LET $param = 5").unwrap();
75		assert_eq!(format!("{}", query), "LET $param = 5;");
76
77		let query = parse("LET $param: number = 5").unwrap();
78		assert_eq!(format!("{}", query), "LET $param: number = 5;");
79	}
80}