surrealdb_core/sql/statements/
update.rs

1use crate::ctx::Context;
2use crate::dbs::{Iterator, Options, Statement};
3use crate::doc::CursorDoc;
4use crate::err::Error;
5use crate::idx::planner::{QueryPlanner, RecordStrategy, StatementContext};
6use crate::sql::{Cond, Data, Explain, Output, Timeout, Value, Values, With};
7
8use reblessive::tree::Stk;
9use revision::revisioned;
10use serde::{Deserialize, Serialize};
11use std::fmt;
12
13#[revisioned(revision = 3)]
14#[derive(Clone, Debug, Default, Eq, PartialEq, PartialOrd, Serialize, Deserialize, Hash)]
15#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
16#[non_exhaustive]
17pub struct UpdateStatement {
18	#[revision(start = 2)]
19	pub only: bool,
20	pub what: Values,
21	#[revision(start = 3)]
22	pub with: Option<With>,
23	pub data: Option<Data>,
24	pub cond: Option<Cond>,
25	pub output: Option<Output>,
26	pub timeout: Option<Timeout>,
27	pub parallel: bool,
28	#[revision(start = 3)]
29	pub explain: Option<Explain>,
30}
31
32impl UpdateStatement {
33	/// Check if we require a writeable transaction
34	pub(crate) fn writeable(&self) -> bool {
35		true
36	}
37	/// Process this type returning a computed simple Value
38	pub(crate) async fn compute(
39		&self,
40		stk: &mut Stk,
41		ctx: &Context,
42		opt: &Options,
43		doc: Option<&CursorDoc>,
44	) -> Result<Value, Error> {
45		// Valid options?
46		opt.valid_for_db()?;
47		// Create a new iterator
48		let mut i = Iterator::new();
49		// Assign the statement
50		let stm = Statement::from(self);
51		// Ensure futures are stored
52		let opt = &opt.new_with_futures(false);
53		// Check if there is a timeout
54		let ctx = stm.setup_timeout(ctx)?;
55		// Get a query planner
56		let mut planner = QueryPlanner::new();
57		let stm_ctx = StatementContext::new(&ctx, opt, &stm)?;
58		// Loop over the update targets
59		for w in self.what.0.iter() {
60			let v = w.compute(stk, &ctx, opt, doc).await?;
61			i.prepare(stk, &mut planner, &stm_ctx, v).await.map_err(|e| match e {
62				Error::InvalidStatementTarget {
63					value: v,
64				} => Error::UpdateStatement {
65					value: v,
66				},
67				e => e,
68			})?;
69		}
70		// Attach the query planner to the context
71		let ctx = stm.setup_query_planner(planner, ctx);
72		// Process the statement
73		let res = i.output(stk, &ctx, opt, &stm, RecordStrategy::KeysAndValues).await?;
74		// Catch statement timeout
75		if ctx.is_timedout()? {
76			return Err(Error::QueryTimedout);
77		}
78		// Output the results
79		match res {
80			// This is a single record result
81			Value::Array(mut a) if self.only => match a.len() {
82				// There was exactly one result
83				1 => Ok(a.remove(0)),
84				// There were no results
85				_ => Err(Error::SingleOnlyOutput),
86			},
87			// This is standard query result
88			v => Ok(v),
89		}
90	}
91}
92
93impl fmt::Display for UpdateStatement {
94	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
95		write!(f, "UPDATE")?;
96		if self.only {
97			f.write_str(" ONLY")?
98		}
99		write!(f, " {}", self.what)?;
100		if let Some(ref v) = self.with {
101			write!(f, " {v}")?
102		}
103		if let Some(ref v) = self.data {
104			write!(f, " {v}")?
105		}
106		if let Some(ref v) = self.cond {
107			write!(f, " {v}")?
108		}
109		if let Some(ref v) = self.output {
110			write!(f, " {v}")?
111		}
112		if let Some(ref v) = self.timeout {
113			write!(f, " {v}")?
114		}
115		if self.parallel {
116			f.write_str(" PARALLEL")?
117		}
118		if let Some(ref v) = self.explain {
119			write!(f, " {v}")?
120		}
121		Ok(())
122	}
123}