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

surrealdb_core/sql/statements/
upsert.rs
