1use super::*;
2
3use super::query::literal_to_sql;
4
5fn resolve_expr_value(expr: &IRExpr, params: &ParamMap) -> Result<Literal> {
9 match expr {
10 IRExpr::Literal(lit) => Ok(lit.clone()),
11 IRExpr::Param(name) => params
12 .get(name)
13 .cloned()
14 .ok_or_else(|| OmniError::manifest(format!("parameter '{}' not provided", name))),
15 other => Err(OmniError::manifest(format!(
16 "unsupported expression in mutation: {:?}",
17 other
18 ))),
19 }
20}
21
22fn literal_to_typed_array(
24 lit: &Literal,
25 data_type: &DataType,
26 num_rows: usize,
27) -> Result<ArrayRef> {
28 Ok(match (lit, data_type) {
29 (Literal::Null, _) => arrow_array::new_null_array(data_type, num_rows),
30 (Literal::String(s), DataType::Utf8) => {
31 Arc::new(StringArray::from(vec![s.as_str(); num_rows])) as ArrayRef
32 }
33 (Literal::Integer(n), DataType::Int32) => {
34 Arc::new(Int32Array::from(vec![*n as i32; num_rows]))
35 }
36 (Literal::Integer(n), DataType::Int64) => Arc::new(Int64Array::from(vec![*n; num_rows])),
37 (Literal::Integer(n), DataType::UInt32) => {
38 Arc::new(UInt32Array::from(vec![*n as u32; num_rows]))
39 }
40 (Literal::Integer(n), DataType::UInt64) => {
41 Arc::new(UInt64Array::from(vec![*n as u64; num_rows]))
42 }
43 (Literal::Float(f), DataType::Float32) => {
44 Arc::new(Float32Array::from(vec![*f as f32; num_rows]))
45 }
46 (Literal::Float(f), DataType::Float64) => Arc::new(Float64Array::from(vec![*f; num_rows])),
47 (Literal::Bool(b), DataType::Boolean) => Arc::new(BooleanArray::from(vec![*b; num_rows])),
48 (Literal::Date(s), DataType::Date32) => {
49 let days = crate::loader::parse_date32_literal(s)?;
50 Arc::new(Date32Array::from(vec![days; num_rows]))
51 }
52 (Literal::DateTime(s), DataType::Date64) => Arc::new(Date64Array::from(vec![
53 crate::loader::parse_date64_literal(s)?;
54 num_rows
55 ])),
56 (Literal::List(items), DataType::List(field)) => {
57 typed_list_literal_to_array(items, field.data_type(), num_rows)?
58 }
59 (Literal::List(items), DataType::FixedSizeList(field, dim))
60 if field.data_type() == &DataType::Float32 =>
61 {
62 if items.len() != *dim as usize {
63 return Err(OmniError::manifest(format!(
64 "vector property expects {} dimensions, got {}",
65 dim,
66 items.len()
67 )));
68 }
69 let mut builder = FixedSizeListBuilder::with_capacity(
70 Float32Builder::with_capacity(num_rows * (*dim as usize)),
71 *dim,
72 num_rows,
73 )
74 .with_field(field.clone());
75 for _ in 0..num_rows {
76 for item in items {
77 match item {
78 Literal::Integer(value) => builder.values().append_value(*value as f32),
79 Literal::Float(value) => builder.values().append_value(*value as f32),
80 _ => {
81 return Err(OmniError::manifest(
82 "vector elements must be numeric".to_string(),
83 ));
84 }
85 }
86 }
87 builder.append(true);
88 }
89 Arc::new(builder.finish())
90 }
91 _ => {
92 return Err(OmniError::manifest(format!(
93 "cannot convert {:?} to {:?}",
94 lit, data_type
95 )));
96 }
97 })
98}
99
100fn typed_list_literal_to_array(
101 items: &[Literal],
102 item_type: &DataType,
103 num_rows: usize,
104) -> Result<ArrayRef> {
105 match item_type {
106 DataType::Utf8 => {
107 let mut builder = ListBuilder::new(StringBuilder::new());
108 for _ in 0..num_rows {
109 for item in items {
110 match item {
111 Literal::String(value) => builder.values().append_value(value),
112 _ => builder.values().append_null(),
113 }
114 }
115 builder.append(true);
116 }
117 Ok(Arc::new(builder.finish()))
118 }
119 DataType::Boolean => {
120 let mut builder = ListBuilder::new(BooleanBuilder::new());
121 for _ in 0..num_rows {
122 for item in items {
123 match item {
124 Literal::Bool(value) => builder.values().append_value(*value),
125 _ => builder.values().append_null(),
126 }
127 }
128 builder.append(true);
129 }
130 Ok(Arc::new(builder.finish()))
131 }
132 DataType::Int32 => {
133 let mut builder = ListBuilder::new(Int32Builder::new());
134 for _ in 0..num_rows {
135 for item in items {
136 match item {
137 Literal::Integer(value) => {
138 let value = i32::try_from(*value).map_err(|_| {
139 OmniError::manifest(format!(
140 "list value {} exceeds Int32 range",
141 value
142 ))
143 })?;
144 builder.values().append_value(value);
145 }
146 _ => builder.values().append_null(),
147 }
148 }
149 builder.append(true);
150 }
151 Ok(Arc::new(builder.finish()))
152 }
153 DataType::Int64 => {
154 let mut builder = ListBuilder::new(Int64Builder::new());
155 for _ in 0..num_rows {
156 for item in items {
157 match item {
158 Literal::Integer(value) => builder.values().append_value(*value),
159 _ => builder.values().append_null(),
160 }
161 }
162 builder.append(true);
163 }
164 Ok(Arc::new(builder.finish()))
165 }
166 DataType::UInt32 => {
167 let mut builder = ListBuilder::new(UInt32Builder::new());
168 for _ in 0..num_rows {
169 for item in items {
170 match item {
171 Literal::Integer(value) => {
172 let value = u32::try_from(*value).map_err(|_| {
173 OmniError::manifest(format!(
174 "list value {} exceeds UInt32 range",
175 value
176 ))
177 })?;
178 builder.values().append_value(value);
179 }
180 _ => builder.values().append_null(),
181 }
182 }
183 builder.append(true);
184 }
185 Ok(Arc::new(builder.finish()))
186 }
187 DataType::UInt64 => {
188 let mut builder = ListBuilder::new(UInt64Builder::new());
189 for _ in 0..num_rows {
190 for item in items {
191 match item {
192 Literal::Integer(value) => {
193 let value = u64::try_from(*value).map_err(|_| {
194 OmniError::manifest(format!(
195 "list value {} exceeds UInt64 range",
196 value
197 ))
198 })?;
199 builder.values().append_value(value);
200 }
201 _ => builder.values().append_null(),
202 }
203 }
204 builder.append(true);
205 }
206 Ok(Arc::new(builder.finish()))
207 }
208 DataType::Float32 => {
209 let mut builder = ListBuilder::new(Float32Builder::new());
210 for _ in 0..num_rows {
211 for item in items {
212 match item {
213 Literal::Integer(value) => builder.values().append_value(*value as f32),
214 Literal::Float(value) => builder.values().append_value(*value as f32),
215 _ => builder.values().append_null(),
216 }
217 }
218 builder.append(true);
219 }
220 Ok(Arc::new(builder.finish()))
221 }
222 DataType::Float64 => {
223 let mut builder = ListBuilder::new(Float64Builder::new());
224 for _ in 0..num_rows {
225 for item in items {
226 match item {
227 Literal::Integer(value) => builder.values().append_value(*value as f64),
228 Literal::Float(value) => builder.values().append_value(*value),
229 _ => builder.values().append_null(),
230 }
231 }
232 builder.append(true);
233 }
234 Ok(Arc::new(builder.finish()))
235 }
236 DataType::Date32 => {
237 let mut builder = ListBuilder::new(Date32Builder::new());
238 for _ in 0..num_rows {
239 for item in items {
240 match item {
241 Literal::Date(value) => builder
242 .values()
243 .append_value(crate::loader::parse_date32_literal(value)?),
244 _ => builder.values().append_null(),
245 }
246 }
247 builder.append(true);
248 }
249 Ok(Arc::new(builder.finish()))
250 }
251 DataType::Date64 => {
252 let mut builder = ListBuilder::new(Date64Builder::new());
253 for _ in 0..num_rows {
254 for item in items {
255 match item {
256 Literal::DateTime(value) => builder
257 .values()
258 .append_value(crate::loader::parse_date64_literal(value)?),
259 _ => builder.values().append_null(),
260 }
261 }
262 builder.append(true);
263 }
264 Ok(Arc::new(builder.finish()))
265 }
266 other => Err(OmniError::manifest(format!(
267 "cannot convert list literal to {:?}",
268 other
269 ))),
270 }
271}
272
273fn build_blob_array_from_value(value: &str) -> Result<ArrayRef> {
275 let mut builder = BlobArrayBuilder::new(1);
276 crate::loader::append_blob_value(&mut builder, value)?;
277 builder
278 .finish()
279 .map_err(|e| OmniError::Lance(e.to_string()))
280}
281
282fn build_null_blob_array() -> Result<ArrayRef> {
284 let mut builder = BlobArrayBuilder::new(1);
285 builder
286 .push_null()
287 .map_err(|e| OmniError::Lance(e.to_string()))?;
288 builder
289 .finish()
290 .map_err(|e| OmniError::Lance(e.to_string()))
291}
292
293fn build_insert_batch(
295 schema: &SchemaRef,
296 id: &str,
297 assignments: &HashMap<String, Literal>,
298 blob_properties: &HashSet<String>,
299) -> Result<RecordBatch> {
300 let mut columns: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
301
302 for field in schema.fields() {
303 if field.name() == "id" {
304 columns.push(Arc::new(StringArray::from(vec![id])));
305 } else if blob_properties.contains(field.name()) {
306 if let Some(Literal::String(uri)) = assignments.get(field.name()) {
307 columns.push(build_blob_array_from_value(uri)?);
308 } else if field.is_nullable() {
309 columns.push(build_null_blob_array()?);
310 } else {
311 return Err(OmniError::manifest(format!(
312 "missing required blob property '{}'",
313 field.name()
314 )));
315 }
316 } else if field.name() == "src" {
317 let lit = assignments.get("from").ok_or_else(|| {
318 OmniError::manifest("missing required edge endpoint 'from'".to_string())
319 })?;
320 columns.push(literal_to_typed_array(lit, field.data_type(), 1)?);
321 } else if field.name() == "dst" {
322 let lit = assignments.get("to").ok_or_else(|| {
323 OmniError::manifest("missing required edge endpoint 'to'".to_string())
324 })?;
325 columns.push(literal_to_typed_array(lit, field.data_type(), 1)?);
326 } else if let Some(lit) = assignments.get(field.name()) {
327 columns.push(literal_to_typed_array(lit, field.data_type(), 1)?);
328 } else if field.is_nullable() {
329 columns.push(arrow_array::new_null_array(field.data_type(), 1));
330 } else {
331 return Err(OmniError::manifest(format!(
332 "missing required property '{}'",
333 field.name()
334 )));
335 }
336 }
337
338 RecordBatch::try_new(schema.clone(), columns).map_err(|e| OmniError::Lance(e.to_string()))
339}
340
341async fn validate_edge_insert_endpoints(
342 db: &Omnigraph,
343 staging: &MutationStaging,
344 branch: Option<&str>,
345 edge_name: &str,
346 assignments: &HashMap<String, Literal>,
347) -> Result<()> {
348 let catalog = db.catalog();
349 let edge_type = catalog
350 .edge_types
351 .get(edge_name)
352 .ok_or_else(|| OmniError::manifest(format!("unknown edge type '{}'", edge_name)))?;
353 let from = match assignments.get("from") {
354 Some(Literal::String(value)) => value.as_str(),
355 Some(other) => {
356 return Err(OmniError::manifest(format!(
357 "edge {} from endpoint must be a string id, got {}",
358 edge_name,
359 literal_to_sql(other)
360 )));
361 }
362 None => {
363 return Err(OmniError::manifest(format!(
364 "edge {} missing 'from' endpoint",
365 edge_name
366 )));
367 }
368 };
369 let to = match assignments.get("to") {
370 Some(Literal::String(value)) => value.as_str(),
371 Some(other) => {
372 return Err(OmniError::manifest(format!(
373 "edge {} to endpoint must be a string id, got {}",
374 edge_name,
375 literal_to_sql(other)
376 )));
377 }
378 None => {
379 return Err(OmniError::manifest(format!(
380 "edge {} missing 'to' endpoint",
381 edge_name
382 )));
383 }
384 };
385
386 ensure_node_id_exists(db, staging, branch, &edge_type.from_type, from, "src").await?;
387 ensure_node_id_exists(db, staging, branch, &edge_type.to_type, to, "dst").await?;
388 Ok(())
389}
390
391fn pending_batches_contain_id(batches: &[RecordBatch], id: &str) -> bool {
395 for batch in batches {
396 let Some(col) = batch.column_by_name("id") else {
397 continue;
398 };
399 let Some(arr) = col.as_any().downcast_ref::<StringArray>() else {
400 continue;
401 };
402 for i in 0..arr.len() {
403 if arr.is_valid(i) && arr.value(i) == id {
404 return true;
405 }
406 }
407 }
408 false
409}
410
411async fn ensure_node_id_exists(
412 db: &Omnigraph,
413 staging: &MutationStaging,
414 branch: Option<&str>,
415 node_type: &str,
416 id: &str,
417 label: &str,
418) -> Result<()> {
419 let table_key = format!("node:{}", node_type);
420
421 let pending = staging.pending_batches(&table_key);
425 if pending_batches_contain_id(pending, id) {
426 return Ok(());
427 }
428
429 let filter = format!("id = '{}'", id.replace('\'', "''"));
430 let snapshot = db.snapshot_for_branch(branch).await?;
431 let ds = db
432 .storage()
433 .open_snapshot_at_table(&snapshot, &table_key)
434 .await?;
435 let exists = db.storage().count_rows(&ds, Some(filter)).await? > 0;
436
437 if exists {
438 Ok(())
439 } else {
440 Err(OmniError::manifest(format!(
441 "{} '{}' not found in {}",
442 label, id, node_type
443 )))
444 }
445}
446
447fn predicate_to_sql(
449 predicate: &IRMutationPredicate,
450 params: &ParamMap,
451 is_edge: bool,
452) -> Result<String> {
453 let column = if is_edge {
454 match predicate.property.as_str() {
455 "from" => "src".to_string(),
456 "to" => "dst".to_string(),
457 other => other.to_string(),
458 }
459 } else {
460 predicate.property.clone()
461 };
462
463 let value = resolve_expr_value(&predicate.value, params)?;
464 let value_sql = literal_to_sql(&value);
465
466 let op = match predicate.op {
467 CompOp::Eq => "=",
468 CompOp::Ne => "!=",
469 CompOp::Gt => ">",
470 CompOp::Lt => "<",
471 CompOp::Ge => ">=",
472 CompOp::Le => "<=",
473 CompOp::Contains => {
474 return Err(OmniError::manifest(
475 "contains predicate not supported in mutations".to_string(),
476 ));
477 }
478 };
479
480 Ok(format!("{} {} {}", column, op, value_sql))
481}
482
483fn apply_assignments(
501 full_schema: &SchemaRef,
502 batch: &RecordBatch,
503 assignments: &HashMap<String, Literal>,
504 blob_properties: &HashSet<String>,
505) -> Result<RecordBatch> {
506 let mut columns: Vec<ArrayRef> = Vec::with_capacity(full_schema.fields().len());
507 let mut out_fields: Vec<Field> = Vec::with_capacity(full_schema.fields().len());
508
509 for field in full_schema.fields().iter() {
510 if blob_properties.contains(field.name()) {
511 if let Some(Literal::String(uri)) = assignments.get(field.name()) {
512 let mut builder = BlobArrayBuilder::new(batch.num_rows());
514 for _ in 0..batch.num_rows() {
515 crate::loader::append_blob_value(&mut builder, uri)?;
516 }
517 let blob_field = lance::blob::blob_field(field.name(), true);
518 out_fields.push(blob_field);
519 columns.push(
520 builder
521 .finish()
522 .map_err(|e| OmniError::Lance(e.to_string()))?,
523 );
524 } else if let Some(col) = batch.column_by_name(field.name()) {
525 let blob_field = lance::blob::blob_field(field.name(), field.is_nullable());
529 out_fields.push(blob_field);
530 columns.push(col.clone());
531 }
532 } else if let Some(lit) = assignments.get(field.name()) {
536 out_fields.push(field.as_ref().clone());
537 columns.push(literal_to_typed_array(
538 lit,
539 field.data_type(),
540 batch.num_rows(),
541 )?);
542 } else {
543 let col = batch.column_by_name(field.name()).ok_or_else(|| {
544 OmniError::Lance(format!(
545 "column '{}' not found in scan result",
546 field.name()
547 ))
548 })?;
549 out_fields.push(field.as_ref().clone());
550 columns.push(col.clone());
551 }
552 }
553
554 RecordBatch::try_new(Arc::new(Schema::new(out_fields)), columns)
555 .map_err(|e| OmniError::Lance(e.to_string()))
556}
557
558use super::staging::{MutationStaging, PendingMode};
561
562async fn open_table_for_mutation(
599 db: &Omnigraph,
600 staging: &mut MutationStaging,
601 branch: Option<&str>,
602 table_key: &str,
603 op_kind: crate::db::MutationOpKind,
604) -> Result<(SnapshotHandle, String, Option<String>)> {
605 if let Some(prior) = staging.inline_committed.get(table_key) {
606 let path = staging.paths.get(table_key).ok_or_else(|| {
607 OmniError::manifest_internal(format!(
608 "open_table_for_mutation: inline_committed[{}] without paths entry",
609 table_key
610 ))
611 })?;
612 let ds = db
613 .reopen_for_mutation(
614 table_key,
615 &path.full_path,
616 path.table_branch.as_deref(),
617 prior.table_version,
618 op_kind,
619 )
620 .await?;
621 return Ok((ds, path.full_path.clone(), path.table_branch.clone()));
622 }
623 let (ds, full_path, table_branch) = db
624 .open_for_mutation_on_branch(branch, table_key, op_kind)
625 .await?;
626 let expected_version = ds.version();
627 staging.ensure_path(
628 table_key,
629 full_path.clone(),
630 table_branch.clone(),
631 expected_version,
632 op_kind,
633 );
634 Ok((ds, full_path, table_branch))
635}
636
637fn enforce_no_mixed_destructive_constructive(
649 ir: &omnigraph_compiler::ir::MutationIR,
650) -> Result<()> {
651 let mut has_constructive = false;
652 let mut has_delete = false;
653 for op in &ir.ops {
654 match op {
655 MutationOpIR::Insert { .. } | MutationOpIR::Update { .. } => {
656 has_constructive = true;
657 }
658 MutationOpIR::Delete { .. } => {
659 has_delete = true;
660 }
661 }
662 }
663 if has_constructive && has_delete {
664 return Err(OmniError::manifest(format!(
665 "mutation '{}' on the same query mixes inserts/updates and deletes; \
666 split into separate mutations: (1) inserts and updates, then (2) deletes. \
667 This restriction lifts when Lance exposes a two-phase delete API \
668 (tracked: lance-format/lance#6658).",
669 ir.name
670 )));
671 }
672 Ok(())
673}
674
675impl Omnigraph {
676 pub async fn mutate(
677 &self,
678 branch: &str,
679 query_source: &str,
680 query_name: &str,
681 params: &ParamMap,
682 ) -> Result<MutationResult> {
683 self.mutate_as(branch, query_source, query_name, params, None)
684 .await
685 }
686
687 pub async fn mutate_as(
688 &self,
689 branch: &str,
690 query_source: &str,
691 query_name: &str,
692 params: &ParamMap,
693 actor_id: Option<&str>,
694 ) -> Result<MutationResult> {
695 self.enforce(
701 omnigraph_policy::PolicyAction::Change,
702 &omnigraph_policy::ResourceScope::Branch(branch.to_string()),
703 actor_id,
704 )?;
705 self.mutate_with_current_actor(branch, query_source, query_name, params, actor_id)
706 .await
707 }
708
709 async fn mutate_with_current_actor(
710 &self,
711 branch: &str,
712 query_source: &str,
713 query_name: &str,
714 params: &ParamMap,
715 actor_id: Option<&str>,
716 ) -> Result<MutationResult> {
717 self.ensure_schema_state_valid().await?;
718 self.heal_pending_recovery_sidecars().await?;
726 let requested = Self::normalize_branch_name(branch)?;
727 if let Some(name) = requested.as_deref() {
732 crate::db::ensure_public_branch_ref(name, "mutate")?;
733 }
734 let resolved_params = enrich_mutation_params(params)?;
735
736 let mut staging = MutationStaging::default();
743
744 let ir = self.lower_named_mutation(query_source, query_name)?;
748
749 let fork_queue_guards: Option<(
756 Vec<(String, Option<String>)>,
757 Vec<tokio::sync::OwnedMutexGuard<()>>,
758 )> = if let Some(active) = requested.as_deref() {
759 let snapshot = self.snapshot_for_branch(Some(active)).await?;
760 let touched: Vec<(String, Option<String>)> = self
761 .touched_table_keys(&ir)
762 .into_iter()
763 .map(|k| (k, Some(active.to_string())))
764 .collect();
765 let needs_fork = touched.iter().any(|(table_key, _)| {
766 snapshot
767 .entry(table_key)
768 .map(|e| e.table_branch.as_deref() != Some(active))
769 .unwrap_or(false)
770 });
771 if needs_fork {
772 let guards = self.write_queue().acquire_many(&touched).await;
773 Some((touched, guards))
774 } else {
775 None
776 }
777 } else {
778 None
779 };
780
781 let exec_result = self
782 .execute_named_mutation(&ir, &resolved_params, requested.as_deref(), &mut staging)
783 .await;
784
785 match exec_result {
786 Err(e) => Err(e),
787 Ok(total) if staging.is_empty() => Ok(total),
788 Ok(total) => {
789 let staged = staging.stage_all(self, requested.as_deref()).await?;
790 let (updates, expected_versions, sidecar_handle, _queue_guards) = staged
797 .commit_all(
798 self,
799 requested.as_deref(),
800 crate::db::manifest::SidecarKind::Mutation,
801 actor_id,
802 fork_queue_guards,
803 )
804 .await?;
805 crate::failpoints::maybe_fail("mutation.post_finalize_pre_publisher")?;
816 self.commit_updates_on_branch_with_expected(
817 requested.as_deref(),
818 &updates,
819 &expected_versions,
820 actor_id,
821 )
822 .await?;
823 if let Some(handle) = sidecar_handle {
829 if let Err(err) =
838 crate::db::manifest::delete_sidecar(&handle, self.storage_adapter()).await
839 {
840 tracing::warn!(
841 error = %err,
842 operation_id = handle.operation_id.as_str(),
843 "recovery sidecar cleanup failed; the next open's recovery sweep will resolve it"
844 );
845 }
846 }
847 Ok(total)
848 }
849 }
850 }
851
852 fn lower_named_mutation(
861 &self,
862 query_source: &str,
863 query_name: &str,
864 ) -> Result<omnigraph_compiler::ir::MutationIR> {
865 let query_decl = omnigraph_compiler::find_named_query(query_source, query_name)
866 .map_err(|e| OmniError::manifest(e.to_string()))?;
867
868 let checked = typecheck_query_decl(&self.catalog(), &query_decl)?;
869 match checked {
870 CheckedQuery::Mutation(_) => {}
871 CheckedQuery::Read(_) => {
872 return Err(OmniError::manifest(
873 "mutation execution called on a read query; use query instead".to_string(),
874 ));
875 }
876 }
877
878 let ir = lower_mutation_query(&query_decl)?;
879 enforce_no_mixed_destructive_constructive(&ir)?;
881 Ok(ir)
882 }
883
884 fn touched_table_keys(&self, ir: &omnigraph_compiler::ir::MutationIR) -> Vec<String> {
899 use omnigraph_compiler::ir::MutationOpIR;
900 let catalog = self.catalog();
901 let mut keys: Vec<String> = Vec::new();
902 for op in &ir.ops {
903 let type_name = match op {
904 MutationOpIR::Insert { type_name, .. }
905 | MutationOpIR::Update { type_name, .. }
906 | MutationOpIR::Delete { type_name, .. } => type_name,
907 };
908 if catalog.node_types.contains_key(type_name) {
909 keys.push(format!("node:{type_name}"));
910 if matches!(op, MutationOpIR::Delete { .. }) {
914 for (edge_name, edge_type) in &catalog.edge_types {
915 if edge_type.from_type == *type_name || edge_type.to_type == *type_name {
916 keys.push(format!("edge:{edge_name}"));
917 }
918 }
919 }
920 } else if catalog.edge_types.contains_key(type_name) {
921 keys.push(format!("edge:{type_name}"));
922 }
923 }
924 keys.sort();
925 keys.dedup();
926 keys
927 }
928
929 async fn execute_named_mutation(
930 &self,
931 ir: &omnigraph_compiler::ir::MutationIR,
932 params: &ParamMap,
933 branch: Option<&str>,
934 staging: &mut MutationStaging,
935 ) -> Result<MutationResult> {
936 let mut total = MutationResult::default();
937 for op in &ir.ops {
938 let result = match op {
939 MutationOpIR::Insert {
940 type_name,
941 assignments,
942 } => {
943 self.execute_insert(type_name, assignments, params, branch, staging)
944 .await?
945 }
946 MutationOpIR::Update {
947 type_name,
948 assignments,
949 predicate,
950 } => {
951 self.execute_update(type_name, assignments, predicate, params, branch, staging)
952 .await?
953 }
954 MutationOpIR::Delete {
955 type_name,
956 predicate,
957 } => {
958 self.execute_delete(type_name, predicate, params, branch, staging)
959 .await?
960 }
961 };
962 total.affected_nodes += result.affected_nodes;
963 total.affected_edges += result.affected_edges;
964 }
965 Ok(total)
966 }
967
968 async fn execute_insert(
969 &self,
970 type_name: &str,
971 assignments: &[IRAssignment],
972 params: &ParamMap,
973 branch: Option<&str>,
974 staging: &mut MutationStaging,
975 ) -> Result<MutationResult> {
976 let mut resolved: HashMap<String, Literal> = HashMap::new();
977 for a in assignments {
978 resolved.insert(a.property.clone(), resolve_expr_value(&a.value, params)?);
979 }
980
981 let is_node = self.catalog().node_types.contains_key(type_name);
982 let is_edge = self.catalog().edge_types.contains_key(type_name);
983
984 if is_node {
985 let node_type = &self.catalog().node_types[type_name];
986 let schema = node_type.arrow_schema.clone();
987 let blob_props = node_type.blob_properties.clone();
988 let id = if let Some(key_prop) = node_type.key_property() {
989 match resolved.get(key_prop) {
990 Some(Literal::String(s)) => s.clone(),
991 Some(other) => literal_to_sql(other).trim_matches('\'').to_string(),
992 None => {
993 return Err(OmniError::manifest(format!(
994 "insert missing @key property '{}'",
995 key_prop
996 )));
997 }
998 }
999 } else {
1000 ulid::Ulid::new().to_string()
1001 };
1002
1003 let batch = build_insert_batch(&schema, &id, &resolved, &blob_props)?;
1004 crate::loader::validate_value_constraints(&batch, node_type)?;
1005 crate::loader::validate_enum_constraints(&batch, &node_type.properties, type_name)?;
1006 let unique_groups = crate::loader::unique_constraint_groups_for_node(node_type);
1007 if !unique_groups.is_empty() {
1008 crate::loader::enforce_unique_constraints_intra_batch(
1009 &batch,
1010 type_name,
1011 &unique_groups,
1012 )?;
1013 }
1014 let has_key = node_type.key_property().is_some();
1015 let table_key = format!("node:{}", type_name);
1016 let insert_kind = if has_key {
1018 crate::db::MutationOpKind::Merge
1019 } else {
1020 crate::db::MutationOpKind::Insert
1021 };
1022 let (_ds, _full_path, _table_branch) =
1023 open_table_for_mutation(self, staging, branch, &table_key, insert_kind).await?;
1024 let mode = if has_key {
1028 PendingMode::Merge
1029 } else {
1030 PendingMode::Append
1031 };
1032 staging.append_batch(&table_key, schema, mode, batch)?;
1033
1034 Ok(MutationResult {
1035 affected_nodes: 1,
1036 affected_edges: 0,
1037 })
1038 } else if is_edge {
1039 let edge_type = &self.catalog().edge_types[type_name];
1040 let schema = edge_type.arrow_schema.clone();
1041 let blob_props = edge_type.blob_properties.clone();
1042 let id = ulid::Ulid::new().to_string();
1043
1044 let batch = build_insert_batch(&schema, &id, &resolved, &blob_props)?;
1045 validate_edge_insert_endpoints(self, staging, branch, type_name, &resolved).await?;
1046 crate::loader::validate_enum_constraints(&batch, &edge_type.properties, type_name)?;
1047 let unique_groups = crate::loader::unique_constraint_groups_for_edge(edge_type);
1048 if !unique_groups.is_empty() {
1049 crate::loader::enforce_unique_constraints_intra_batch(
1050 &batch,
1051 type_name,
1052 &unique_groups,
1053 )?;
1054 }
1055 let table_key = format!("edge:{}", type_name);
1056 let (ds, _full_path, _table_branch) = open_table_for_mutation(
1058 self,
1059 staging,
1060 branch,
1061 &table_key,
1062 crate::db::MutationOpKind::Insert,
1063 )
1064 .await?;
1065 staging.append_batch(&table_key, schema, PendingMode::Append, batch.clone())?;
1068
1069 validate_edge_cardinality_with_pending(self, &ds, staging, &table_key, edge_type)
1074 .await?;
1075
1076 self.invalidate_graph_index().await;
1077
1078 Ok(MutationResult {
1079 affected_nodes: 0,
1080 affected_edges: 1,
1081 })
1082 } else {
1083 Err(OmniError::manifest(format!("unknown type '{}'", type_name)))
1084 }
1085 }
1086
1087 async fn execute_update(
1088 &self,
1089 type_name: &str,
1090 assignments: &[IRAssignment],
1091 predicate: &IRMutationPredicate,
1092 params: &ParamMap,
1093 branch: Option<&str>,
1094 staging: &mut MutationStaging,
1095 ) -> Result<MutationResult> {
1096 if !self.catalog().node_types.contains_key(type_name) {
1098 return Err(OmniError::manifest(format!(
1099 "update is only supported for node types, not '{}'",
1100 type_name
1101 )));
1102 }
1103
1104 if let Some(key_prop) = self.catalog().node_types[type_name].key_property() {
1106 if assignments.iter().any(|a| a.property == key_prop) {
1107 return Err(OmniError::manifest(format!(
1108 "cannot update @key property '{}' — delete and re-insert instead",
1109 key_prop
1110 )));
1111 }
1112 }
1113
1114 let pred_sql = predicate_to_sql(predicate, params, false)?;
1115 let schema = self.catalog().node_types[type_name].arrow_schema.clone();
1116 let blob_props = self.catalog().node_types[type_name].blob_properties.clone();
1117
1118 let table_key = format!("node:{}", type_name);
1119 let (ds, _full_path, _table_branch) = open_table_for_mutation(
1120 self,
1121 staging,
1122 branch,
1123 &table_key,
1124 crate::db::MutationOpKind::Update,
1125 )
1126 .await?;
1127
1128 let non_blob_cols: Vec<&str> = schema
1141 .fields()
1142 .iter()
1143 .filter(|f| !blob_props.contains(f.name()))
1144 .map(|f| f.name().as_str())
1145 .collect();
1146 let projection: Option<&[&str]> =
1147 (!blob_props.is_empty()).then_some(non_blob_cols.as_slice());
1148 let pending_batches = staging.pending_batches(&table_key);
1149 let pending_schema = staging.pending_schema(&table_key);
1150 let batches = self
1157 .storage()
1158 .scan_with_pending(
1159 &ds,
1160 pending_batches,
1161 pending_schema,
1162 projection,
1163 Some(&pred_sql),
1164 Some("id"),
1165 )
1166 .await?;
1167
1168 if batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0) {
1169 return Ok(MutationResult {
1170 affected_nodes: 0,
1171 affected_edges: 0,
1172 });
1173 }
1174
1175 let matched = concat_match_batches_to_schema(&schema, &blob_props, batches)?;
1184
1185 let affected_count = matched.num_rows();
1186
1187 let mut resolved: HashMap<String, Literal> = HashMap::new();
1188 for a in assignments {
1189 resolved.insert(a.property.clone(), resolve_expr_value(&a.value, params)?);
1190 }
1191 let updated = apply_assignments(&schema, &matched, &resolved, &blob_props)?;
1192 let node_type = &self.catalog().node_types[type_name];
1193 crate::loader::validate_value_constraints(&updated, node_type)?;
1194 crate::loader::validate_enum_constraints(&updated, &node_type.properties, type_name)?;
1195 let unique_groups = crate::loader::unique_constraint_groups_for_node(node_type);
1196 if !unique_groups.is_empty() {
1197 crate::loader::enforce_unique_constraints_intra_batch(
1198 &updated,
1199 type_name,
1200 &unique_groups,
1201 )?;
1202 }
1203
1204 let updated_schema = updated.schema();
1210 staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?;
1211
1212 Ok(MutationResult {
1213 affected_nodes: affected_count,
1214 affected_edges: 0,
1215 })
1216 }
1217
1218 async fn execute_delete(
1219 &self,
1220 type_name: &str,
1221 predicate: &IRMutationPredicate,
1222 params: &ParamMap,
1223 branch: Option<&str>,
1224 staging: &mut MutationStaging,
1225 ) -> Result<MutationResult> {
1226 let is_node = self.catalog().node_types.contains_key(type_name);
1227 if is_node {
1228 self.execute_delete_node(type_name, predicate, params, branch, staging)
1229 .await
1230 } else {
1231 self.execute_delete_edge(type_name, predicate, params, branch, staging)
1232 .await
1233 }
1234 }
1235
1236 async fn execute_delete_node(
1237 &self,
1238 type_name: &str,
1239 predicate: &IRMutationPredicate,
1240 params: &ParamMap,
1241 branch: Option<&str>,
1242 staging: &mut MutationStaging,
1243 ) -> Result<MutationResult> {
1244 let pred_sql = predicate_to_sql(predicate, params, false)?;
1245
1246 let table_key = format!("node:{}", type_name);
1247 let (ds, full_path, table_branch) = open_table_for_mutation(
1248 self,
1249 staging,
1250 branch,
1251 &table_key,
1252 crate::db::MutationOpKind::Delete,
1253 )
1254 .await?;
1255 let initial_version = ds.version();
1256
1257 let batches = self
1261 .storage()
1262 .scan(&ds, Some(&["id"]), Some(&pred_sql), None)
1263 .await?;
1264
1265 let deleted_ids: Vec<String> = batches
1266 .iter()
1267 .flat_map(|batch| {
1268 let ids = batch
1269 .column(0)
1270 .as_any()
1271 .downcast_ref::<StringArray>()
1272 .unwrap();
1273 (0..ids.len())
1274 .map(|i| ids.value(i).to_string())
1275 .collect::<Vec<_>>()
1276 })
1277 .collect();
1278
1279 if deleted_ids.is_empty() {
1280 return Ok(MutationResult {
1281 affected_nodes: 0,
1282 affected_edges: 0,
1283 });
1284 }
1285
1286 let affected_nodes = deleted_ids.len();
1287
1288 let ds = self
1294 .reopen_for_mutation(
1295 &table_key,
1296 &full_path,
1297 table_branch.as_deref(),
1298 initial_version,
1299 crate::db::MutationOpKind::Delete,
1300 )
1301 .await?;
1302 crate::failpoints::maybe_fail("mutation.delete_node_pre_primary_delete")?;
1303 let (_new_ds, delete_state) = self
1304 .storage_inline_residual()
1305 .delete_where(&full_path, ds, &pred_sql)
1306 .await?;
1307
1308 staging.record_inline(crate::db::SubTableUpdate {
1309 table_key: table_key.clone(),
1310 table_version: delete_state.version,
1311 table_branch: table_branch.clone(),
1312 row_count: delete_state.row_count,
1313 version_metadata: delete_state.version_metadata,
1314 });
1315
1316 let mut affected_edges = 0usize;
1317 let escaped: Vec<String> = deleted_ids
1318 .iter()
1319 .map(|id| format!("'{}'", id.replace('\'', "''")))
1320 .collect();
1321 let id_list = escaped.join(", ");
1322
1323 let edge_info: Vec<(String, String, String)> = self
1324 .catalog()
1325 .edge_types
1326 .iter()
1327 .map(|(name, et)| (name.clone(), et.from_type.clone(), et.to_type.clone()))
1328 .collect();
1329
1330 for (edge_name, from_type, to_type) in &edge_info {
1331 let mut cascade_filters = Vec::new();
1332 if from_type == type_name {
1333 cascade_filters.push(format!("src IN ({})", id_list));
1334 }
1335 if to_type == type_name {
1336 cascade_filters.push(format!("dst IN ({})", id_list));
1337 }
1338 if cascade_filters.is_empty() {
1339 continue;
1340 }
1341
1342 let edge_table_key = format!("edge:{}", edge_name);
1343 let cascade_filter = cascade_filters.join(" OR ");
1344 let (edge_ds, edge_full_path, edge_table_branch) = open_table_for_mutation(
1345 self,
1346 staging,
1347 branch,
1348 &edge_table_key,
1349 crate::db::MutationOpKind::Delete,
1350 )
1351 .await?;
1352
1353 let (_new_edge_ds, edge_delete) = self
1354 .storage_inline_residual()
1355 .delete_where(&edge_full_path, edge_ds, &cascade_filter)
1356 .await?;
1357
1358 affected_edges += edge_delete.deleted_rows;
1359
1360 if edge_delete.deleted_rows > 0 {
1361 staging.record_inline(crate::db::SubTableUpdate {
1362 table_key: edge_table_key,
1363 table_version: edge_delete.version,
1364 table_branch: edge_table_branch,
1365 row_count: edge_delete.row_count,
1366 version_metadata: edge_delete.version_metadata,
1367 });
1368 }
1369 }
1370
1371 if affected_edges > 0 {
1372 self.invalidate_graph_index().await;
1373 }
1374
1375 Ok(MutationResult {
1376 affected_nodes,
1377 affected_edges,
1378 })
1379 }
1380
1381 async fn execute_delete_edge(
1382 &self,
1383 type_name: &str,
1384 predicate: &IRMutationPredicate,
1385 params: &ParamMap,
1386 branch: Option<&str>,
1387 staging: &mut MutationStaging,
1388 ) -> Result<MutationResult> {
1389 let pred_sql = predicate_to_sql(predicate, params, true)?;
1390
1391 let table_key = format!("edge:{}", type_name);
1392 let (ds, full_path, table_branch) = open_table_for_mutation(
1393 self,
1394 staging,
1395 branch,
1396 &table_key,
1397 crate::db::MutationOpKind::Delete,
1398 )
1399 .await?;
1400
1401 let (_new_ds, delete_state) = self
1402 .storage_inline_residual()
1403 .delete_where(&full_path, ds, &pred_sql)
1404 .await?;
1405 let affected = delete_state.deleted_rows;
1406
1407 if affected > 0 {
1408 staging.record_inline(crate::db::SubTableUpdate {
1409 table_key,
1410 table_version: delete_state.version,
1411 table_branch,
1412 row_count: delete_state.row_count,
1413 version_metadata: delete_state.version_metadata,
1414 });
1415 self.invalidate_graph_index().await;
1416 }
1417
1418 Ok(MutationResult {
1419 affected_nodes: 0,
1420 affected_edges: affected,
1421 })
1422 }
1423}
1424
1425fn concat_match_batches_to_schema(
1431 _schema: &SchemaRef,
1432 _blob_properties: &HashSet<String>,
1433 batches: Vec<RecordBatch>,
1434) -> Result<RecordBatch> {
1435 if batches.len() == 1 {
1436 return Ok(batches.into_iter().next().unwrap());
1437 }
1438 let common = batches[0].schema();
1439 arrow_select::concat::concat_batches(&common, &batches).map_err(|e| {
1440 OmniError::Lance(format!(
1441 "scan_with_pending returned batches with mismatched schemas \
1442 across the committed/pending boundary; this typically indicates \
1443 a blob-column shape mismatch between the committed table and a \
1444 prior in-query insert/update. Split blob-touching mutations \
1445 into separate queries. ({})",
1446 e
1447 ))
1448 })
1449}
1450
1451async fn validate_edge_cardinality_with_pending(
1456 db: &Omnigraph,
1457 committed_ds: &SnapshotHandle,
1458 staging: &MutationStaging,
1459 table_key: &str,
1460 edge_type: &omnigraph_compiler::catalog::EdgeType,
1461) -> Result<()> {
1462 if edge_type.cardinality.is_default() {
1463 return Ok(());
1464 }
1465 let counts =
1466 super::staging::count_src_per_edge(db, committed_ds, table_key, staging, None).await?;
1467 super::staging::enforce_cardinality_bounds(edge_type, &counts)
1468}
1469
1470fn enrich_mutation_params(params: &ParamMap) -> Result<ParamMap> {
1471 let mut resolved = params.clone();
1472 if !resolved.contains_key(NOW_PARAM_NAME) {
1473 let now = OffsetDateTime::now_utc()
1474 .format(&Rfc3339)
1475 .map_err(|e| OmniError::manifest(format!("failed to format now(): {}", e)))?;
1476 resolved.insert(NOW_PARAM_NAME.to_string(), Literal::DateTime(now));
1477 }
1478 Ok(resolved)
1479}