surrealdb_core/sql/statements/
update.rs

1use crate::ctx::{Context, MutableContext};
2use crate::dbs::{Iterator, Options, Statement};
3use crate::doc::CursorDoc;
4use crate::err::Error;
5use crate::idx::planner::RecordStrategy;
6use crate::sql::{Cond, Data, Output, Timeout, Value, Values};
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 UpdateStatement {
18	#[revision(start = 2)]
19	pub only: bool,
20	pub what: Values,
21	pub data: Option<Data>,
22	pub cond: Option<Cond>,
23	pub output: Option<Output>,
24	pub timeout: Option<Timeout>,
25	pub parallel: bool,
26}
27
28impl UpdateStatement {
29	/// Check if we require a writeable transaction
30	pub(crate) fn writeable(&self) -> bool {
31		true
32	}
33	/// Process this type returning a computed simple Value
34	pub(crate) async fn compute(
35		&self,
36		stk: &mut Stk,
37		ctx: &Context,
38		opt: &Options,
39		doc: Option<&CursorDoc>,
40	) -> Result<Value, Error> {
41		// Valid options?
42		opt.valid_for_db()?;
43		// Create a new iterator
44		let mut i = Iterator::new();
45		// Assign the statement
46		let stm = Statement::from(self);
47		// Ensure futures are stored
48		let opt = &opt.new_with_futures(false);
49		// Check if there is a timeout
50		let ctx = match self.timeout.as_ref() {
51			Some(timeout) => {
52				let mut ctx = MutableContext::new(ctx);
53				ctx.add_timeout(*timeout.0)?;
54				ctx.freeze()
55			}
56			None => ctx.clone(),
57		};
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(&stm, v).map_err(|e| match e {
62				Error::InvalidStatementTarget {
63					value: v,
64				} => Error::UpdateStatement {
65					value: v,
66				},
67				e => e,
68			})?;
69		}
70		// Process the statement
71		let res = i.output(stk, &ctx, opt, &stm, RecordStrategy::KeysAndValues).await?;
72		// Catch statement timeout
73		if ctx.is_timedout() {
74			return Err(Error::QueryTimedout);
75		}
76		// Output the results
77		match res {
78			// This is a single record result
79			Value::Array(mut a) if self.only => match a.len() {
80				// There was exactly one result
81				1 => Ok(a.remove(0)),
82				// There were no results
83				_ => Err(Error::SingleOnlyOutput),
84			},
85			// This is standard query result
86			v => Ok(v),
87		}
88	}
89}
90
91impl fmt::Display for UpdateStatement {
92	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
93		write!(f, "UPDATE")?;
94		if self.only {
95			f.write_str(" ONLY")?
96		}
97		write!(f, " {}", self.what)?;
98		if let Some(ref v) = self.data {
99			write!(f, " {v}")?
100		}
101		if let Some(ref v) = self.cond {
102			write!(f, " {v}")?
103		}
104		if let Some(ref v) = self.output {
105			write!(f, " {v}")?
106		}
107		if let Some(ref v) = self.timeout {
108			write!(f, " {v}")?
109		}
110		if self.parallel {
111			f.write_str(" PARALLEL")?
112		}
113		Ok(())
114	}
115}