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 = snapshot.open(&table_key).await?;
432 let exists = ds
433 .count_rows(Some(filter))
434 .await
435 .map_err(|e| OmniError::Lance(e.to_string()))?
436 > 0;
437
438 if exists {
439 Ok(())
440 } else {
441 Err(OmniError::manifest(format!(
442 "{} '{}' not found in {}",
443 label, id, node_type
444 )))
445 }
446}
447
448fn predicate_to_sql(
450 predicate: &IRMutationPredicate,
451 params: &ParamMap,
452 is_edge: bool,
453) -> Result<String> {
454 let column = if is_edge {
455 match predicate.property.as_str() {
456 "from" => "src".to_string(),
457 "to" => "dst".to_string(),
458 other => other.to_string(),
459 }
460 } else {
461 predicate.property.clone()
462 };
463
464 let value = resolve_expr_value(&predicate.value, params)?;
465 let value_sql = literal_to_sql(&value);
466
467 let op = match predicate.op {
468 CompOp::Eq => "=",
469 CompOp::Ne => "!=",
470 CompOp::Gt => ">",
471 CompOp::Lt => "<",
472 CompOp::Ge => ">=",
473 CompOp::Le => "<=",
474 CompOp::Contains => {
475 return Err(OmniError::manifest(
476 "contains predicate not supported in mutations".to_string(),
477 ));
478 }
479 };
480
481 Ok(format!("{} {} {}", column, op, value_sql))
482}
483
484fn apply_assignments(
502 full_schema: &SchemaRef,
503 batch: &RecordBatch,
504 assignments: &HashMap<String, Literal>,
505 blob_properties: &HashSet<String>,
506) -> Result<RecordBatch> {
507 let mut columns: Vec<ArrayRef> = Vec::with_capacity(full_schema.fields().len());
508 let mut out_fields: Vec<Field> = Vec::with_capacity(full_schema.fields().len());
509
510 for field in full_schema.fields().iter() {
511 if blob_properties.contains(field.name()) {
512 if let Some(Literal::String(uri)) = assignments.get(field.name()) {
513 let mut builder = BlobArrayBuilder::new(batch.num_rows());
515 for _ in 0..batch.num_rows() {
516 crate::loader::append_blob_value(&mut builder, uri)?;
517 }
518 let blob_field = lance::blob::blob_field(field.name(), true);
519 out_fields.push(blob_field);
520 columns.push(
521 builder
522 .finish()
523 .map_err(|e| OmniError::Lance(e.to_string()))?,
524 );
525 } else if let Some(col) = batch.column_by_name(field.name()) {
526 let blob_field = lance::blob::blob_field(field.name(), field.is_nullable());
530 out_fields.push(blob_field);
531 columns.push(col.clone());
532 }
533 } else if let Some(lit) = assignments.get(field.name()) {
537 out_fields.push(field.as_ref().clone());
538 columns.push(literal_to_typed_array(
539 lit,
540 field.data_type(),
541 batch.num_rows(),
542 )?);
543 } else {
544 let col = batch.column_by_name(field.name()).ok_or_else(|| {
545 OmniError::Lance(format!(
546 "column '{}' not found in scan result",
547 field.name()
548 ))
549 })?;
550 out_fields.push(field.as_ref().clone());
551 columns.push(col.clone());
552 }
553 }
554
555 RecordBatch::try_new(Arc::new(Schema::new(out_fields)), columns)
556 .map_err(|e| OmniError::Lance(e.to_string()))
557}
558
559use super::staging::{MutationStaging, PendingMode};
562
563async fn open_table_for_mutation(
600 db: &Omnigraph,
601 staging: &mut MutationStaging,
602 branch: Option<&str>,
603 table_key: &str,
604 op_kind: crate::db::MutationOpKind,
605) -> Result<(Dataset, String, Option<String>)> {
606 if let Some(prior) = staging.inline_committed.get(table_key) {
607 let path = staging.paths.get(table_key).ok_or_else(|| {
608 OmniError::manifest_internal(format!(
609 "open_table_for_mutation: inline_committed[{}] without paths entry",
610 table_key
611 ))
612 })?;
613 let ds = db
614 .reopen_for_mutation(
615 table_key,
616 &path.full_path,
617 path.table_branch.as_deref(),
618 prior.table_version,
619 op_kind,
620 )
621 .await?;
622 return Ok((ds, path.full_path.clone(), path.table_branch.clone()));
623 }
624 let (ds, full_path, table_branch) = db
625 .open_for_mutation_on_branch(branch, table_key, op_kind)
626 .await?;
627 let expected_version = ds.version().version;
628 staging.ensure_path(
629 table_key,
630 full_path.clone(),
631 table_branch.clone(),
632 expected_version,
633 op_kind,
634 );
635 Ok((ds, full_path, table_branch))
636}
637
638fn enforce_no_mixed_destructive_constructive(
650 ir: &omnigraph_compiler::ir::MutationIR,
651) -> Result<()> {
652 let mut has_constructive = false;
653 let mut has_delete = false;
654 for op in &ir.ops {
655 match op {
656 MutationOpIR::Insert { .. } | MutationOpIR::Update { .. } => {
657 has_constructive = true;
658 }
659 MutationOpIR::Delete { .. } => {
660 has_delete = true;
661 }
662 }
663 }
664 if has_constructive && has_delete {
665 return Err(OmniError::manifest(format!(
666 "mutation '{}' on the same query mixes inserts/updates and deletes; \
667 split into separate mutations: (1) inserts and updates, then (2) deletes. \
668 This restriction lifts when Lance exposes a two-phase delete API \
669 (tracked: lance-format/lance#6658).",
670 ir.name
671 )));
672 }
673 Ok(())
674}
675
676impl Omnigraph {
677 pub async fn mutate(
678 &self,
679 branch: &str,
680 query_source: &str,
681 query_name: &str,
682 params: &ParamMap,
683 ) -> Result<MutationResult> {
684 self.mutate_as(branch, query_source, query_name, params, None)
685 .await
686 }
687
688 pub async fn mutate_as(
689 &self,
690 branch: &str,
691 query_source: &str,
692 query_name: &str,
693 params: &ParamMap,
694 actor_id: Option<&str>,
695 ) -> Result<MutationResult> {
696 self.enforce(
702 omnigraph_policy::PolicyAction::Change,
703 &omnigraph_policy::ResourceScope::Branch(branch.to_string()),
704 actor_id,
705 )?;
706 self.mutate_with_current_actor(branch, query_source, query_name, params, actor_id)
707 .await
708 }
709
710 async fn mutate_with_current_actor(
711 &self,
712 branch: &str,
713 query_source: &str,
714 query_name: &str,
715 params: &ParamMap,
716 actor_id: Option<&str>,
717 ) -> Result<MutationResult> {
718 self.ensure_schema_state_valid().await?;
719 let requested = Self::normalize_branch_name(branch)?;
720 if let Some(name) = requested.as_deref() {
725 crate::db::ensure_public_branch_ref(name, "mutate")?;
726 }
727 let resolved_params = enrich_mutation_params(params)?;
728
729 let mut staging = MutationStaging::default();
736
737 let exec_result = self
738 .execute_named_mutation(
739 query_source,
740 query_name,
741 &resolved_params,
742 requested.as_deref(),
743 &mut staging,
744 )
745 .await;
746
747 match exec_result {
748 Err(e) => Err(e),
749 Ok(total) if staging.is_empty() => Ok(total),
750 Ok(total) => {
751 let staged = staging.stage_all(self, requested.as_deref()).await?;
752 let (updates, expected_versions, sidecar_handle, _queue_guards) = staged
759 .commit_all(
760 self,
761 requested.as_deref(),
762 crate::db::manifest::SidecarKind::Mutation,
763 actor_id,
764 )
765 .await?;
766 crate::failpoints::maybe_fail("mutation.post_finalize_pre_publisher")?;
777 self.commit_updates_on_branch_with_expected(
778 requested.as_deref(),
779 &updates,
780 &expected_versions,
781 actor_id,
782 )
783 .await?;
784 if let Some(handle) = sidecar_handle {
790 if let Err(err) =
799 crate::db::manifest::delete_sidecar(&handle, self.storage_adapter()).await
800 {
801 tracing::warn!(
802 error = %err,
803 operation_id = handle.operation_id.as_str(),
804 "recovery sidecar cleanup failed; the next open's recovery sweep will resolve it"
805 );
806 }
807 }
808 Ok(total)
809 }
810 }
811 }
812
813 async fn execute_named_mutation(
814 &self,
815 query_source: &str,
816 query_name: &str,
817 params: &ParamMap,
818 branch: Option<&str>,
819 staging: &mut MutationStaging,
820 ) -> Result<MutationResult> {
821 let query_decl = omnigraph_compiler::find_named_query(query_source, query_name)
822 .map_err(|e| OmniError::manifest(e.to_string()))?;
823
824 let checked = typecheck_query_decl(&self.catalog(), &query_decl)?;
825 match checked {
826 CheckedQuery::Mutation(_) => {}
827 CheckedQuery::Read(_) => {
828 return Err(OmniError::manifest(
829 "mutation execution called on a read query; use query instead".to_string(),
830 ));
831 }
832 }
833
834 let ir = lower_mutation_query(&query_decl)?;
835 enforce_no_mixed_destructive_constructive(&ir)?;
837
838 let mut total = MutationResult::default();
839 for op in &ir.ops {
840 let result = match op {
841 MutationOpIR::Insert {
842 type_name,
843 assignments,
844 } => {
845 self.execute_insert(type_name, assignments, params, branch, staging)
846 .await?
847 }
848 MutationOpIR::Update {
849 type_name,
850 assignments,
851 predicate,
852 } => {
853 self.execute_update(type_name, assignments, predicate, params, branch, staging)
854 .await?
855 }
856 MutationOpIR::Delete {
857 type_name,
858 predicate,
859 } => {
860 self.execute_delete(type_name, predicate, params, branch, staging)
861 .await?
862 }
863 };
864 total.affected_nodes += result.affected_nodes;
865 total.affected_edges += result.affected_edges;
866 }
867 Ok(total)
868 }
869
870 async fn execute_insert(
871 &self,
872 type_name: &str,
873 assignments: &[IRAssignment],
874 params: &ParamMap,
875 branch: Option<&str>,
876 staging: &mut MutationStaging,
877 ) -> Result<MutationResult> {
878 let mut resolved: HashMap<String, Literal> = HashMap::new();
879 for a in assignments {
880 resolved.insert(a.property.clone(), resolve_expr_value(&a.value, params)?);
881 }
882
883 let is_node = self.catalog().node_types.contains_key(type_name);
884 let is_edge = self.catalog().edge_types.contains_key(type_name);
885
886 if is_node {
887 let node_type = &self.catalog().node_types[type_name];
888 let schema = node_type.arrow_schema.clone();
889 let blob_props = node_type.blob_properties.clone();
890 let id = if let Some(key_prop) = node_type.key_property() {
891 match resolved.get(key_prop) {
892 Some(Literal::String(s)) => s.clone(),
893 Some(other) => literal_to_sql(other).trim_matches('\'').to_string(),
894 None => {
895 return Err(OmniError::manifest(format!(
896 "insert missing @key property '{}'",
897 key_prop
898 )));
899 }
900 }
901 } else {
902 ulid::Ulid::new().to_string()
903 };
904
905 let batch = build_insert_batch(&schema, &id, &resolved, &blob_props)?;
906 crate::loader::validate_value_constraints(&batch, node_type)?;
907 crate::loader::validate_enum_constraints(&batch, &node_type.properties, type_name)?;
908 let unique_props = crate::loader::unique_property_names_for_node(node_type);
909 if !unique_props.is_empty() {
910 crate::loader::enforce_unique_constraints_intra_batch(
911 &batch,
912 type_name,
913 &unique_props,
914 )?;
915 }
916 let has_key = node_type.key_property().is_some();
917 let table_key = format!("node:{}", type_name);
918 let insert_kind = if has_key {
920 crate::db::MutationOpKind::Merge
921 } else {
922 crate::db::MutationOpKind::Insert
923 };
924 let (_ds, _full_path, _table_branch) =
925 open_table_for_mutation(self, staging, branch, &table_key, insert_kind).await?;
926 let mode = if has_key {
930 PendingMode::Merge
931 } else {
932 PendingMode::Append
933 };
934 staging.append_batch(&table_key, schema, mode, batch)?;
935
936 Ok(MutationResult {
937 affected_nodes: 1,
938 affected_edges: 0,
939 })
940 } else if is_edge {
941 let edge_type = &self.catalog().edge_types[type_name];
942 let schema = edge_type.arrow_schema.clone();
943 let blob_props = edge_type.blob_properties.clone();
944 let id = ulid::Ulid::new().to_string();
945
946 let batch = build_insert_batch(&schema, &id, &resolved, &blob_props)?;
947 validate_edge_insert_endpoints(self, staging, branch, type_name, &resolved).await?;
948 crate::loader::validate_enum_constraints(&batch, &edge_type.properties, type_name)?;
949 let unique_props = crate::loader::unique_property_names_for_edge(edge_type);
950 if !unique_props.is_empty() {
951 crate::loader::enforce_unique_constraints_intra_batch(
952 &batch,
953 type_name,
954 &unique_props,
955 )?;
956 }
957 let table_key = format!("edge:{}", type_name);
958 let (ds, _full_path, _table_branch) = open_table_for_mutation(
960 self,
961 staging,
962 branch,
963 &table_key,
964 crate::db::MutationOpKind::Insert,
965 )
966 .await?;
967 staging.append_batch(&table_key, schema, PendingMode::Append, batch.clone())?;
970
971 validate_edge_cardinality_with_pending(self, &ds, staging, &table_key, edge_type)
976 .await?;
977
978 self.invalidate_graph_index().await;
979
980 Ok(MutationResult {
981 affected_nodes: 0,
982 affected_edges: 1,
983 })
984 } else {
985 Err(OmniError::manifest(format!("unknown type '{}'", type_name)))
986 }
987 }
988
989 async fn execute_update(
990 &self,
991 type_name: &str,
992 assignments: &[IRAssignment],
993 predicate: &IRMutationPredicate,
994 params: &ParamMap,
995 branch: Option<&str>,
996 staging: &mut MutationStaging,
997 ) -> Result<MutationResult> {
998 if !self.catalog().node_types.contains_key(type_name) {
1000 return Err(OmniError::manifest(format!(
1001 "update is only supported for node types, not '{}'",
1002 type_name
1003 )));
1004 }
1005
1006 if let Some(key_prop) = self.catalog().node_types[type_name].key_property() {
1008 if assignments.iter().any(|a| a.property == key_prop) {
1009 return Err(OmniError::manifest(format!(
1010 "cannot update @key property '{}' — delete and re-insert instead",
1011 key_prop
1012 )));
1013 }
1014 }
1015
1016 let pred_sql = predicate_to_sql(predicate, params, false)?;
1017 let schema = self.catalog().node_types[type_name].arrow_schema.clone();
1018 let blob_props = self.catalog().node_types[type_name].blob_properties.clone();
1019
1020 let table_key = format!("node:{}", type_name);
1021 let (ds, _full_path, _table_branch) = open_table_for_mutation(
1022 self,
1023 staging,
1024 branch,
1025 &table_key,
1026 crate::db::MutationOpKind::Update,
1027 )
1028 .await?;
1029
1030 let non_blob_cols: Vec<&str> = schema
1043 .fields()
1044 .iter()
1045 .filter(|f| !blob_props.contains(f.name()))
1046 .map(|f| f.name().as_str())
1047 .collect();
1048 let projection: Option<&[&str]> =
1049 (!blob_props.is_empty()).then_some(non_blob_cols.as_slice());
1050 let pending_batches = staging.pending_batches(&table_key);
1051 let pending_schema = staging.pending_schema(&table_key);
1052 let batches = self
1059 .table_store()
1060 .scan_with_pending(
1061 &ds,
1062 pending_batches,
1063 pending_schema,
1064 projection,
1065 Some(&pred_sql),
1066 Some("id"),
1067 )
1068 .await?;
1069
1070 if batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0) {
1071 return Ok(MutationResult {
1072 affected_nodes: 0,
1073 affected_edges: 0,
1074 });
1075 }
1076
1077 let matched = concat_match_batches_to_schema(&schema, &blob_props, batches)?;
1086
1087 let affected_count = matched.num_rows();
1088
1089 let mut resolved: HashMap<String, Literal> = HashMap::new();
1090 for a in assignments {
1091 resolved.insert(a.property.clone(), resolve_expr_value(&a.value, params)?);
1092 }
1093 let updated = apply_assignments(&schema, &matched, &resolved, &blob_props)?;
1094 let node_type = &self.catalog().node_types[type_name];
1095 crate::loader::validate_value_constraints(&updated, node_type)?;
1096 crate::loader::validate_enum_constraints(&updated, &node_type.properties, type_name)?;
1097 let unique_props = crate::loader::unique_property_names_for_node(node_type);
1098 if !unique_props.is_empty() {
1099 crate::loader::enforce_unique_constraints_intra_batch(
1100 &updated,
1101 type_name,
1102 &unique_props,
1103 )?;
1104 }
1105
1106 let updated_schema = updated.schema();
1112 staging.append_batch(&table_key, updated_schema, PendingMode::Merge, updated)?;
1113
1114 Ok(MutationResult {
1115 affected_nodes: affected_count,
1116 affected_edges: 0,
1117 })
1118 }
1119
1120 async fn execute_delete(
1121 &self,
1122 type_name: &str,
1123 predicate: &IRMutationPredicate,
1124 params: &ParamMap,
1125 branch: Option<&str>,
1126 staging: &mut MutationStaging,
1127 ) -> Result<MutationResult> {
1128 let is_node = self.catalog().node_types.contains_key(type_name);
1129 if is_node {
1130 self.execute_delete_node(type_name, predicate, params, branch, staging)
1131 .await
1132 } else {
1133 self.execute_delete_edge(type_name, predicate, params, branch, staging)
1134 .await
1135 }
1136 }
1137
1138 async fn execute_delete_node(
1139 &self,
1140 type_name: &str,
1141 predicate: &IRMutationPredicate,
1142 params: &ParamMap,
1143 branch: Option<&str>,
1144 staging: &mut MutationStaging,
1145 ) -> Result<MutationResult> {
1146 let pred_sql = predicate_to_sql(predicate, params, false)?;
1147
1148 let table_key = format!("node:{}", type_name);
1149 let (ds, full_path, table_branch) = open_table_for_mutation(
1150 self,
1151 staging,
1152 branch,
1153 &table_key,
1154 crate::db::MutationOpKind::Delete,
1155 )
1156 .await?;
1157 let initial_version = ds.version().version;
1158
1159 let batches = self
1163 .table_store()
1164 .scan(&ds, Some(&["id"]), Some(&pred_sql), None)
1165 .await?;
1166
1167 let deleted_ids: Vec<String> = batches
1168 .iter()
1169 .flat_map(|batch| {
1170 let ids = batch
1171 .column(0)
1172 .as_any()
1173 .downcast_ref::<StringArray>()
1174 .unwrap();
1175 (0..ids.len())
1176 .map(|i| ids.value(i).to_string())
1177 .collect::<Vec<_>>()
1178 })
1179 .collect();
1180
1181 if deleted_ids.is_empty() {
1182 return Ok(MutationResult {
1183 affected_nodes: 0,
1184 affected_edges: 0,
1185 });
1186 }
1187
1188 let affected_nodes = deleted_ids.len();
1189
1190 let mut ds = self
1196 .reopen_for_mutation(
1197 &table_key,
1198 &full_path,
1199 table_branch.as_deref(),
1200 initial_version,
1201 crate::db::MutationOpKind::Delete,
1202 )
1203 .await?;
1204 crate::failpoints::maybe_fail("mutation.delete_node_pre_primary_delete")?;
1205 let delete_state = self
1206 .table_store()
1207 .delete_where(&full_path, &mut ds, &pred_sql)
1208 .await?;
1209
1210 staging.record_inline(crate::db::SubTableUpdate {
1211 table_key: table_key.clone(),
1212 table_version: delete_state.version,
1213 table_branch: table_branch.clone(),
1214 row_count: delete_state.row_count,
1215 version_metadata: delete_state.version_metadata,
1216 });
1217
1218 let mut affected_edges = 0usize;
1219 let escaped: Vec<String> = deleted_ids
1220 .iter()
1221 .map(|id| format!("'{}'", id.replace('\'', "''")))
1222 .collect();
1223 let id_list = escaped.join(", ");
1224
1225 let edge_info: Vec<(String, String, String)> = self
1226 .catalog()
1227 .edge_types
1228 .iter()
1229 .map(|(name, et)| (name.clone(), et.from_type.clone(), et.to_type.clone()))
1230 .collect();
1231
1232 for (edge_name, from_type, to_type) in &edge_info {
1233 let mut cascade_filters = Vec::new();
1234 if from_type == type_name {
1235 cascade_filters.push(format!("src IN ({})", id_list));
1236 }
1237 if to_type == type_name {
1238 cascade_filters.push(format!("dst IN ({})", id_list));
1239 }
1240 if cascade_filters.is_empty() {
1241 continue;
1242 }
1243
1244 let edge_table_key = format!("edge:{}", edge_name);
1245 let cascade_filter = cascade_filters.join(" OR ");
1246 let (mut edge_ds, edge_full_path, edge_table_branch) = open_table_for_mutation(
1247 self,
1248 staging,
1249 branch,
1250 &edge_table_key,
1251 crate::db::MutationOpKind::Delete,
1252 )
1253 .await?;
1254
1255 let edge_delete = self
1256 .table_store()
1257 .delete_where(&edge_full_path, &mut edge_ds, &cascade_filter)
1258 .await?;
1259
1260 affected_edges += edge_delete.deleted_rows;
1261
1262 if edge_delete.deleted_rows > 0 {
1263 staging.record_inline(crate::db::SubTableUpdate {
1264 table_key: edge_table_key,
1265 table_version: edge_delete.version,
1266 table_branch: edge_table_branch,
1267 row_count: edge_delete.row_count,
1268 version_metadata: edge_delete.version_metadata,
1269 });
1270 }
1271 }
1272
1273 if affected_edges > 0 {
1274 self.invalidate_graph_index().await;
1275 }
1276
1277 Ok(MutationResult {
1278 affected_nodes,
1279 affected_edges,
1280 })
1281 }
1282
1283 async fn execute_delete_edge(
1284 &self,
1285 type_name: &str,
1286 predicate: &IRMutationPredicate,
1287 params: &ParamMap,
1288 branch: Option<&str>,
1289 staging: &mut MutationStaging,
1290 ) -> Result<MutationResult> {
1291 let pred_sql = predicate_to_sql(predicate, params, true)?;
1292
1293 let table_key = format!("edge:{}", type_name);
1294 let (mut ds, full_path, table_branch) = open_table_for_mutation(
1295 self,
1296 staging,
1297 branch,
1298 &table_key,
1299 crate::db::MutationOpKind::Delete,
1300 )
1301 .await?;
1302
1303 let delete_state = self
1304 .table_store()
1305 .delete_where(&full_path, &mut ds, &pred_sql)
1306 .await?;
1307 let affected = delete_state.deleted_rows;
1308
1309 if affected > 0 {
1310 staging.record_inline(crate::db::SubTableUpdate {
1311 table_key,
1312 table_version: delete_state.version,
1313 table_branch,
1314 row_count: delete_state.row_count,
1315 version_metadata: delete_state.version_metadata,
1316 });
1317 self.invalidate_graph_index().await;
1318 }
1319
1320 Ok(MutationResult {
1321 affected_nodes: 0,
1322 affected_edges: affected,
1323 })
1324 }
1325}
1326
1327fn concat_match_batches_to_schema(
1333 _schema: &SchemaRef,
1334 _blob_properties: &HashSet<String>,
1335 batches: Vec<RecordBatch>,
1336) -> Result<RecordBatch> {
1337 if batches.len() == 1 {
1338 return Ok(batches.into_iter().next().unwrap());
1339 }
1340 let common = batches[0].schema();
1341 arrow_select::concat::concat_batches(&common, &batches).map_err(|e| {
1342 OmniError::Lance(format!(
1343 "scan_with_pending returned batches with mismatched schemas \
1344 across the committed/pending boundary; this typically indicates \
1345 a blob-column shape mismatch between the committed table and a \
1346 prior in-query insert/update. Split blob-touching mutations \
1347 into separate queries. ({})",
1348 e
1349 ))
1350 })
1351}
1352
1353async fn validate_edge_cardinality_with_pending(
1358 db: &Omnigraph,
1359 committed_ds: &Dataset,
1360 staging: &MutationStaging,
1361 table_key: &str,
1362 edge_type: &omnigraph_compiler::catalog::EdgeType,
1363) -> Result<()> {
1364 if edge_type.cardinality.is_default() {
1365 return Ok(());
1366 }
1367 let counts =
1368 super::staging::count_src_per_edge(db, committed_ds, table_key, staging, None).await?;
1369 super::staging::enforce_cardinality_bounds(edge_type, &counts)
1370}
1371
1372fn enrich_mutation_params(params: &ParamMap) -> Result<ParamMap> {
1373 let mut resolved = params.clone();
1374 if !resolved.contains_key(NOW_PARAM_NAME) {
1375 let now = OffsetDateTime::now_utc()
1376 .format(&Rfc3339)
1377 .map_err(|e| OmniError::manifest(format!("failed to format now(): {}", e)))?;
1378 resolved.insert(NOW_PARAM_NAME.to_string(), Literal::DateTime(now));
1379 }
1380 Ok(resolved)
1381}