surrealdb_core/sql/statements/
foreach.rs

1use crate::ctx::{Context, MutableContext};
2use crate::dbs::Options;
3use crate::doc::CursorDoc;
4use crate::err::Error;
5use crate::sql::{block::Entry, Block, Param, Value};
6
7use reblessive::tree::Stk;
8use revision::revisioned;
9use serde::{Deserialize, Serialize};
10use std::fmt::{self, Display};
11use std::ops::Deref;
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 ForeachStatement {
18	pub param: Param,
19	pub range: Value,
20	pub block: Block,
21}
22
23enum ForeachIter {
24	Array(std::vec::IntoIter<Value>),
25	Range(std::iter::Map<std::ops::Range<i64>, fn(i64) -> Value>),
26}
27
28impl Iterator for ForeachIter {
29	type Item = Value;
30
31	fn next(&mut self) -> Option<Self::Item> {
32		match self {
33			ForeachIter::Array(iter) => iter.next(),
34			ForeachIter::Range(iter) => iter.next(),
35		}
36	}
37}
38
39impl ForeachStatement {
40	/// Check if we require a writeable transaction
41	pub(crate) fn writeable(&self) -> bool {
42		self.range.writeable() || self.block.writeable()
43	}
44	/// Process this type returning a computed simple Value
45	///
46	/// Was marked recursive
47	pub(crate) async fn compute(
48		&self,
49		stk: &mut Stk,
50		ctx: &Context,
51		opt: &Options,
52		doc: Option<&CursorDoc>,
53	) -> Result<Value, Error> {
54		// Check the loop data
55		let data = self.range.compute(stk, ctx, opt, doc).await?;
56		let iter = match data {
57			Value::Array(arr) => ForeachIter::Array(arr.into_iter()),
58			Value::Range(r) => {
59				let r: std::ops::Range<i64> = r.deref().to_owned().try_into()?;
60				ForeachIter::Range(r.map(Value::from))
61			}
62
63			v => {
64				return Err(Error::InvalidStatementTarget {
65					value: v.to_string(),
66				})
67			}
68		};
69
70		// Loop over the values
71		for v in iter {
72			if ctx.is_timedout()? {
73				return Err(Error::QueryTimedout);
74			}
75			// Duplicate context
76			let ctx = MutableContext::new(ctx).freeze();
77			// Set the current parameter
78			let key = self.param.0.to_raw();
79			let val = stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await?;
80			let mut ctx = MutableContext::unfreeze(ctx)?;
81			ctx.add_value(key, val.into());
82			let mut ctx = ctx.freeze();
83			// Loop over the code block statements
84			for v in self.block.iter() {
85				// Compute each block entry
86				let res = match v {
87					Entry::Set(v) => {
88						let val = stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await?;
89						let mut c = MutableContext::unfreeze(ctx)?;
90						c.add_value(v.name.to_owned(), val.into());
91						ctx = c.freeze();
92						Ok(Value::None)
93					}
94					Entry::Value(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
95					Entry::Break(v) => v.compute(&ctx, opt, doc).await,
96					Entry::Continue(v) => v.compute(&ctx, opt, doc).await,
97					Entry::Foreach(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
98					Entry::Ifelse(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
99					Entry::Select(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
100					Entry::Create(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
101					Entry::Upsert(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
102					Entry::Update(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
103					Entry::Delete(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
104					Entry::Relate(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
105					Entry::Insert(v) => stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await,
106					Entry::Define(v) => v.compute(stk, &ctx, opt, doc).await,
107					Entry::Alter(v) => v.compute(stk, &ctx, opt, doc).await,
108					Entry::Rebuild(v) => v.compute(stk, &ctx, opt, doc).await,
109					Entry::Remove(v) => v.compute(&ctx, opt, doc).await,
110					Entry::Output(v) => {
111						return stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await;
112					}
113					Entry::Throw(v) => {
114						return stk.run(|stk| v.compute(stk, &ctx, opt, doc)).await;
115					}
116				};
117				// Catch any special errors
118				match res {
119					Err(Error::Continue) => break,
120					Err(Error::Break) => return Ok(Value::None),
121					Err(err) => return Err(err),
122					_ => (),
123				};
124			}
125			// Cooperitively yield if the task has been running for too long.
126			yield_now!();
127		}
128		// Ok all good
129		Ok(Value::None)
130	}
131}
132
133impl Display for ForeachStatement {
134	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
135		write!(f, "FOR {} IN {} {}", self.param, self.range, self.block)
136	}
137}
surrealdb_core/sql/statements/foreach.rs

surrealdb_core/sql/statements/
foreach.rs
