use std::sync::Arc;
use futures::{StreamExt, stream};
use surrealdb_types::ToSql;
use crate::catalog::{DatabaseId, NamespaceId};
use crate::ctx::FrozenContext;
use crate::dbs::Options;
use crate::doc::CursorDoc;
use crate::exec::operators::check_cancelled;
use crate::exec::plan_or_compute::{get_legacy_context, legacy_compute};
use crate::exec::{
AccessMode, CardinalityHint, ContextLevel, ExecOperator, ExecutionContext, FlowResult,
OperatorMetrics, ValueBatch, ValueBatchStream, buffer_stream, monitor_stream,
};
use crate::expr::data::Assignment;
use crate::expr::match_plan::{DetachMode, UpdateData};
use crate::expr::statements::{CreateStatement, DeleteStatement, RelateStatement, UpdateStatement};
use crate::expr::{AssignOperator, ControlFlow, Data, Expr, Literal, Output};
use crate::idx::planner::ScanDirection;
use crate::key::graph;
use crate::val::{Object, RecordId, TableName, Value};
#[derive(Debug, Clone)]
pub struct UpdateBinding {
pub(crate) input: Arc<dyn ExecOperator>,
pub(crate) target: String,
pub(crate) data: UpdateData,
pub(crate) metrics: Arc<OperatorMetrics>,
}
impl UpdateBinding {
pub(crate) fn new(input: Arc<dyn ExecOperator>, target: String, data: UpdateData) -> Self {
Self {
input,
target,
data,
metrics: Arc::new(OperatorMetrics::new()),
}
}
}
impl ExecOperator for UpdateBinding {
fn name(&self) -> &'static str {
"UpdateBinding"
}
fn attrs(&self) -> Vec<(String, String)> {
let kind = match &self.data {
UpdateData::Set(_) => "set",
UpdateData::Unset(_) => "unset",
UpdateData::Content(_) => "content",
};
vec![("binding".to_string(), self.target.clone()), ("op".to_string(), kind.to_string())]
}
fn required_context(&self) -> ContextLevel {
ContextLevel::Database.max(self.input.required_context())
}
fn access_mode(&self) -> AccessMode {
AccessMode::ReadWrite
}
fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
vec![&self.input]
}
fn metrics(&self) -> Option<&OperatorMetrics> {
Some(&self.metrics)
}
fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
let input_stream = buffer_stream(
self.input.execute(ctx)?,
self.input.access_mode(),
self.input.cardinality_hint(),
ctx.root().ctx.config.operator_buffer_size,
);
let (opt, frozen) = legacy_handles(ctx)?;
let target = self.target.clone();
let data = self.data.clone();
let ctx = ctx.clone();
let stream = async_stream::try_stream! {
futures::pin_mut!(input_stream);
let mut rows = Vec::new();
while let Some(batch) = input_stream.next().await {
check_cancelled(&ctx)?;
rows.extend(batch?.values);
}
let mut out = Vec::with_capacity(rows.len());
for mut row in rows {
check_cancelled(&ctx)?;
if let Some(rid) = target_record_id(&row, &target) {
let after = apply_update(&data, &rid, &row, &frozen, &opt).await?;
set_binding(&mut row, &target, after);
}
out.push(row);
}
yield ValueBatch { values: out };
};
Ok(monitor_stream(Box::pin(stream), "UpdateBinding", &self.metrics))
}
}
#[derive(Debug, Clone)]
pub struct DeleteBinding {
pub(crate) input: Arc<dyn ExecOperator>,
pub(crate) target: String,
pub(crate) detach: DetachMode,
pub(crate) is_edge: bool,
pub(crate) metrics: Arc<OperatorMetrics>,
}
impl DeleteBinding {
pub(crate) fn new(
input: Arc<dyn ExecOperator>,
target: String,
detach: DetachMode,
is_edge: bool,
) -> Self {
Self {
input,
target,
detach,
is_edge,
metrics: Arc::new(OperatorMetrics::new()),
}
}
}
impl ExecOperator for DeleteBinding {
fn name(&self) -> &'static str {
"DeleteBinding"
}
fn attrs(&self) -> Vec<(String, String)> {
let detach = match self.detach {
DetachMode::Detach => "detach",
DetachMode::NoDetach => "nodetach",
};
vec![("binding".to_string(), self.target.clone()), ("mode".to_string(), detach.to_string())]
}
fn required_context(&self) -> ContextLevel {
ContextLevel::Database.max(self.input.required_context())
}
fn access_mode(&self) -> AccessMode {
AccessMode::ReadWrite
}
fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
vec![&self.input]
}
fn metrics(&self) -> Option<&OperatorMetrics> {
Some(&self.metrics)
}
fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
let input_stream = buffer_stream(
self.input.execute(ctx)?,
self.input.access_mode(),
self.input.cardinality_hint(),
ctx.root().ctx.config.operator_buffer_size,
);
let (opt, frozen) = legacy_handles(ctx)?;
let (ns, db) = {
let db_ctx = ctx.database().map_err(|e| ControlFlow::Err(anyhow::anyhow!(e)))?;
(db_ctx.ns_ctx.ns.namespace_id, db_ctx.db.database_id)
};
let target = self.target.clone();
let detach = self.detach;
let is_edge = self.is_edge;
let ctx = ctx.clone();
let stream = async_stream::try_stream! {
futures::pin_mut!(input_stream);
let mut rows = Vec::new();
while let Some(batch) = input_stream.next().await {
check_cancelled(&ctx)?;
rows.extend(batch?.values);
}
let mut out = Vec::with_capacity(rows.len());
let mut deleted: std::collections::HashSet<RecordId> = std::collections::HashSet::new();
for mut row in rows {
check_cancelled(&ctx)?;
if let Some(rid) = target_record_id(&row, &target) {
if deleted.insert(rid.clone()) {
apply_delete(&rid, detach, is_edge, ns, db, &frozen, &opt).await?;
}
set_binding(&mut row, &target, Value::Null);
null_incident_edges(&mut row, &rid);
}
out.push(row);
}
yield ValueBatch { values: out };
};
Ok(monitor_stream(Box::pin(stream), "DeleteBinding", &self.metrics))
}
}
fn legacy_handles(ctx: &ExecutionContext) -> FlowResult<(Options, FrozenContext)> {
let (opt, frozen) =
get_legacy_context(ctx).map_err(|e| ControlFlow::Err(anyhow::anyhow!(e)))?;
Ok((opt, frozen))
}
fn target_record_id(row: &Value, name: &str) -> Option<RecordId> {
crate::exec::operators::binding_record_id(row, name)
}
fn set_binding(row: &mut Value, name: &str, value: Value) {
if let Value::Object(obj) = row {
obj.insert(name.to_string(), value);
}
}
fn null_incident_edges(row: &mut Value, deleted: &RecordId) {
let Value::Object(obj) = row else {
return;
};
let deleted = Value::RecordId(deleted.clone());
for value in obj.values_mut() {
if let Value::Object(edge) = value
&& (edge.get("in") == Some(&deleted) || edge.get("out") == Some(&deleted))
{
*value = Value::Null;
}
}
}
async fn apply_update(
data: &UpdateData,
rid: &RecordId,
row: &Value,
frozen: &FrozenContext,
opt: &Options,
) -> Result<Value, ControlFlow> {
let update_data = match data {
UpdateData::Set(assignments) => {
let cursor = CursorDoc::new(None, None, row.clone());
let mut assigns = Vec::with_capacity(assignments.len());
for (place, value_expr) in assignments {
let value = legacy_compute(value_expr, frozen, opt, Some(&cursor)).await?;
assigns.push(Assignment {
place: place.clone(),
operator: AssignOperator::Assign,
value: value.into_literal(),
});
}
Data::SetExpression(assigns)
}
UpdateData::Unset(fields) => Data::UnsetExpression(fields.clone()),
UpdateData::Content(object_expr) => {
let value = eval_in_row(object_expr, row, frozen, opt).await?;
Data::ContentExpression(value.into_literal())
}
};
let stmt = UpdateStatement {
what: vec![record_id_expr(rid)],
data: Some(update_data),
output: Some(Output::After),
..Default::default()
};
let result = legacy_compute(&Expr::Update(Box::new(stmt)), frozen, opt, None).await?;
Ok(single_record(result))
}
async fn apply_delete(
rid: &RecordId,
detach: DetachMode,
is_edge: bool,
ns: NamespaceId,
db: DatabaseId,
frozen: &FrozenContext,
opt: &Options,
) -> Result<(), ControlFlow> {
if !is_edge
&& matches!(detach, DetachMode::NoDetach)
&& has_connected_edges(rid, ns, db, frozen).await?
{
return Err(ControlFlow::Err(anyhow::anyhow!(crate::err::Error::InvalidStatement(
format!(
"Cannot DELETE `{}` because it still has connected edges; use `DETACH DELETE` to \
remove the edges as well",
rid.to_sql()
)
))));
}
let stmt = DeleteStatement {
what: vec![record_id_expr(rid)],
..Default::default()
};
legacy_compute(&Expr::Delete(Box::new(stmt)), frozen, opt, None).await?;
Ok(())
}
async fn has_connected_edges(
rid: &RecordId,
ns: NamespaceId,
db: DatabaseId,
frozen: &FrozenContext,
) -> Result<bool, ControlFlow> {
let txn = frozen.tx();
let prefix = graph::prefix(ns, db, &rid.table, &rid.key).map_err(ControlFlow::Err)?;
let suffix = graph::suffix(ns, db, &rid.table, &rid.key).map_err(ControlFlow::Err)?;
let mut cursor = txn
.open_keys_cursor(prefix..suffix, ScanDirection::Forward, 0, None)
.await
.map_err(ControlFlow::Err)?;
let batch = cursor.next_batch(1).await.map_err(ControlFlow::Err)?;
Ok(!batch.is_empty())
}
fn record_id_expr(rid: &RecordId) -> Expr {
Expr::Literal(Literal::RecordId(rid.clone().into_literal()))
}
fn single_record(value: Value) -> Value {
match value {
Value::Array(mut arr) => arr.0.drain(..).next().unwrap_or(Value::Null),
other => other,
}
}
async fn eval_in_row(
expr: &Expr,
row: &Value,
frozen: &FrozenContext,
opt: &Options,
) -> Result<Value, ControlFlow> {
let cursor = CursorDoc::new(None, None, row.clone());
legacy_compute(expr, frozen, opt, Some(&cursor)).await
}
#[derive(Debug, Clone)]
pub(crate) struct InsertNodeOp {
pub(crate) name: String,
pub(crate) table: TableName,
pub(crate) props: Expr,
}
#[derive(Debug, Clone)]
pub(crate) struct InsertEdgeOp {
pub(crate) name: String,
pub(crate) table: TableName,
pub(crate) from: String,
pub(crate) to: String,
pub(crate) props: Expr,
}
#[derive(Debug, Clone)]
pub struct InsertGraph {
pub(crate) input: Arc<dyn ExecOperator>,
pub(crate) nodes: Vec<InsertNodeOp>,
pub(crate) edges: Vec<InsertEdgeOp>,
pub(crate) metrics: Arc<OperatorMetrics>,
}
impl InsertGraph {
pub(crate) fn new(
input: Arc<dyn ExecOperator>,
nodes: Vec<InsertNodeOp>,
edges: Vec<InsertEdgeOp>,
) -> Self {
Self {
input,
nodes,
edges,
metrics: Arc::new(OperatorMetrics::new()),
}
}
}
impl ExecOperator for InsertGraph {
fn name(&self) -> &'static str {
"InsertGraph"
}
fn attrs(&self) -> Vec<(String, String)> {
vec![
("nodes".to_string(), self.nodes.len().to_string()),
("edges".to_string(), self.edges.len().to_string()),
]
}
fn required_context(&self) -> ContextLevel {
ContextLevel::Database.max(self.input.required_context())
}
fn access_mode(&self) -> AccessMode {
AccessMode::ReadWrite
}
fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
vec![&self.input]
}
fn metrics(&self) -> Option<&OperatorMetrics> {
Some(&self.metrics)
}
fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
let input_stream = buffer_stream(
self.input.execute(ctx)?,
self.input.access_mode(),
self.input.cardinality_hint(),
ctx.root().ctx.config.operator_buffer_size,
);
let (opt, frozen) = legacy_handles(ctx)?;
let nodes = self.nodes.clone();
let edges = self.edges.clone();
let ctx = ctx.clone();
let stream = async_stream::try_stream! {
futures::pin_mut!(input_stream);
let mut rows = Vec::new();
while let Some(batch) = input_stream.next().await {
check_cancelled(&ctx)?;
rows.extend(batch?.values);
}
let mut out = Vec::with_capacity(rows.len());
for mut row in rows {
check_cancelled(&ctx)?;
insert_row(&nodes, &edges, &mut row, &frozen, &opt).await?;
out.push(row);
}
yield ValueBatch { values: out };
};
Ok(monitor_stream(Box::pin(stream), "InsertGraph", &self.metrics))
}
}
async fn insert_row(
nodes: &[InsertNodeOp],
edges: &[InsertEdgeOp],
row: &mut Value,
frozen: &FrozenContext,
opt: &Options,
) -> Result<(), ControlFlow> {
for node in nodes {
let props = eval_in_row(&node.props, row, frozen, opt).await?;
let stmt = CreateStatement {
what: vec![Expr::Table(node.table.clone())],
data: Some(Data::ContentExpression(props.into_literal())),
output: Some(Output::After),
..Default::default()
};
let result = legacy_compute(&Expr::Create(Box::new(stmt)), frozen, opt, None).await?;
set_binding(row, &node.name, single_record(result));
}
for edge in edges {
let from = target_record_id(row, &edge.from).ok_or_else(|| endpoint_error(&edge.from))?;
let to = target_record_id(row, &edge.to).ok_or_else(|| endpoint_error(&edge.to))?;
let props = eval_in_row(&edge.props, row, frozen, opt).await?;
let stmt = RelateStatement {
only: false,
or_update: false,
through: Expr::Table(edge.table.clone()),
from: record_id_expr(&from),
to: record_id_expr(&to),
data: Some(Data::ContentExpression(props.into_literal())),
output: Some(Output::After),
timeout: Expr::Literal(Literal::None),
};
let result = legacy_compute(&Expr::Relate(Box::new(stmt)), frozen, opt, None).await?;
set_binding(row, &edge.name, single_record(result));
}
Ok(())
}
fn endpoint_error(name: &str) -> ControlFlow {
ControlFlow::Err(anyhow::anyhow!(crate::err::Error::InvalidStatement(format!(
"INSERT edge endpoint `{name}` did not resolve to a record"
))))
}
#[derive(Debug, Clone)]
pub struct SingleRowScan {
pub(crate) metrics: Arc<OperatorMetrics>,
}
impl SingleRowScan {
pub(crate) fn new() -> Self {
Self {
metrics: Arc::new(OperatorMetrics::new()),
}
}
}
impl Default for SingleRowScan {
fn default() -> Self {
Self::new()
}
}
impl ExecOperator for SingleRowScan {
fn name(&self) -> &'static str {
"SingleRowScan"
}
fn required_context(&self) -> ContextLevel {
ContextLevel::Root
}
fn access_mode(&self) -> AccessMode {
AccessMode::ReadOnly
}
fn cardinality_hint(&self) -> CardinalityHint {
CardinalityHint::Bounded(1)
}
fn metrics(&self) -> Option<&OperatorMetrics> {
Some(&self.metrics)
}
fn execute(&self, _ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
let batch = ValueBatch {
values: vec![Value::Object(Object::default())],
};
let stream = stream::once(async move { Ok(batch) });
Ok(monitor_stream(Box::pin(stream), "SingleRowScan", &self.metrics))
}
}
#[derive(Debug, Clone)]
pub struct DrainSink {
pub(crate) input: Arc<dyn ExecOperator>,
pub(crate) metrics: Arc<OperatorMetrics>,
}
impl DrainSink {
pub(crate) fn new(input: Arc<dyn ExecOperator>) -> Self {
Self {
input,
metrics: Arc::new(OperatorMetrics::new()),
}
}
}
impl ExecOperator for DrainSink {
fn name(&self) -> &'static str {
"DrainSink"
}
fn required_context(&self) -> ContextLevel {
self.input.required_context()
}
fn access_mode(&self) -> AccessMode {
self.input.access_mode()
}
fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
vec![&self.input]
}
fn metrics(&self) -> Option<&OperatorMetrics> {
Some(&self.metrics)
}
fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
let input_stream = buffer_stream(
self.input.execute(ctx)?,
self.input.access_mode(),
self.input.cardinality_hint(),
ctx.root().ctx.config.operator_buffer_size,
);
let ctx = ctx.clone();
let stream = async_stream::try_stream! {
futures::pin_mut!(input_stream);
while let Some(batch) = input_stream.next().await {
check_cancelled(&ctx)?;
let _ = batch?;
}
yield ValueBatch { values: Vec::new() };
};
Ok(monitor_stream(Box::pin(stream), "DrainSink", &self.metrics))
}
}