1use crate::error::FaucetError;
4use serde::{Deserialize, Serialize};
5use serde_json::{Map, Value};
6use std::collections::HashMap;
7
8#[derive(
15 Debug, Clone, Copy, Default, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq,
16)]
17#[serde(rename_all = "snake_case")]
18#[non_exhaustive]
19pub enum WriteMode {
20 #[default]
22 Append,
23 Upsert,
25 Delete,
27 Overwrite,
32}
33
34impl WriteMode {
35 pub fn as_str(&self) -> &'static str {
37 match self {
38 WriteMode::Append => "append",
39 WriteMode::Upsert => "upsert",
40 WriteMode::Delete => "delete",
41 WriteMode::Overwrite => "overwrite",
42 }
43 }
44}
45
46#[derive(Debug, Clone, Default, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)]
49pub struct DeleteMarker {
50 pub field: String,
52 pub values: Vec<String>,
54}
55
56#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
62#[serde(rename_all = "snake_case")]
63#[non_exhaustive]
64pub enum OverwriteScope {
65 Window {
67 column: String,
69 from: Value,
71 to: Value,
73 },
74}
75
76impl OverwriteScope {
77 pub fn column(&self) -> &str {
79 match self {
80 OverwriteScope::Window { column, .. } => column,
81 }
82 }
83
84 pub fn validate(&self) -> Result<(), FaucetError> {
86 match self {
87 OverwriteScope::Window { column, from, to } => {
88 if column.trim().is_empty() {
89 return Err(FaucetError::Config(
90 "overwrite scope: window `column` must not be empty".into(),
91 ));
92 }
93 if from.is_null() || to.is_null() {
94 return Err(FaucetError::Config(
95 "overwrite scope: window `from`/`to` must not be null".into(),
96 ));
97 }
98 Ok(())
99 }
100 }
101 }
102
103 pub fn render_where_literal(&self, quoted_col: &str) -> String {
110 self.render_where_with(quoted_col, sql_literal)
111 }
112
113 pub fn render_where_with(
117 &self,
118 quoted_col: &str,
119 literal: impl Fn(&Value) -> String,
120 ) -> String {
121 match self {
122 OverwriteScope::Window { from, to, .. } => format!(
123 "{quoted_col} >= {} AND {quoted_col} < {}",
124 literal(from),
125 literal(to)
126 ),
127 }
128 }
129}
130
131pub fn sql_literal(v: &Value) -> String {
151 match v {
152 Value::String(s) => format!("'{}'", s.replace('\'', "''")),
153 Value::Number(n) => n.to_string(),
154 Value::Bool(b) => b.to_string(),
155 _ => "NULL".to_owned(),
158 }
159}
160
161#[derive(Debug, Clone, Default, Serialize, Deserialize, schemars::JsonSchema)]
165pub struct WriteSpec {
166 #[serde(default)]
168 pub write_mode: WriteMode,
169 #[serde(default)]
171 pub key: Vec<String>,
172 #[serde(default, skip_serializing_if = "Option::is_none")]
176 pub delete_marker: Option<DeleteMarker>,
177 #[serde(default, skip_serializing_if = "Option::is_none")]
182 pub rollback: Option<crate::rollback::RollbackWriteSpec>,
183}
184
185impl WriteSpec {
186 pub fn rollback_run_id(&self) -> Option<&str> {
188 self.rollback.as_ref().map(|r| r.run_id.as_str())
189 }
190
191 pub fn journals(&self) -> bool {
193 self.rollback.as_ref().is_some_and(|r| r.journal)
194 }
195
196 pub fn keeps_previous(&self) -> bool {
198 self.rollback.as_ref().is_some_and(|r| r.keep_previous)
199 }
200
201 pub fn validate(&self) -> Result<(), FaucetError> {
203 if matches!(self.write_mode, WriteMode::Upsert | WriteMode::Delete) && self.key.is_empty() {
204 return Err(FaucetError::Config(format!(
205 "write_mode: {} requires a non-empty `key`",
206 self.write_mode.as_str()
207 )));
208 }
209 Ok(())
210 }
211
212 pub fn dedups_by_key(&self) -> bool {
217 matches!(self.write_mode, WriteMode::Upsert | WriteMode::Delete) && !self.key.is_empty()
218 }
219
220 pub fn is_overwrite(&self) -> bool {
225 matches!(self.write_mode, WriteMode::Overwrite)
226 }
227}
228
229#[derive(Debug, Clone, PartialEq)]
231pub struct KeyTuple(pub Vec<(String, Value)>);
232
233#[derive(Debug, Default)]
237pub struct WritePlan {
238 pub upserts: Vec<Value>,
240 pub deletes: Vec<KeyTuple>,
242 pub failed: Vec<(usize, String)>,
244}
245
246#[derive(Clone)]
247enum Action {
248 Upsert(Value),
249 Delete(KeyTuple),
250}
251
252pub fn plan_writes(page: &[Value], spec: &WriteSpec) -> WritePlan {
256 debug_assert!(
257 matches!(spec.write_mode, WriteMode::Upsert | WriteMode::Delete),
258 "plan_writes is only for Upsert/Delete — Append and Overwrite are routed separately"
259 );
260 let mut plan = WritePlan::default();
261 let mut index: HashMap<String, usize> = HashMap::new();
262 let mut order: Vec<Action> = Vec::new();
263
264 for (i, rec) in page.iter().enumerate() {
265 let key_tuple = match extract_key(rec, &spec.key) {
266 Ok(k) => k,
267 Err(msg) => {
268 plan.failed.push((i, msg));
269 continue;
270 }
271 };
272 let canon = canonical(&key_tuple);
273
274 let is_delete = match spec.write_mode {
275 WriteMode::Delete => true,
276 WriteMode::Upsert => is_delete_marked(rec, spec.delete_marker.as_ref()),
277 WriteMode::Append | WriteMode::Overwrite => false,
278 };
279
280 let action = if is_delete {
281 Action::Delete(key_tuple)
282 } else {
283 Action::Upsert(strip_marker(rec.clone(), spec.delete_marker.as_ref()))
284 };
285
286 match index.get(&canon) {
287 Some(&slot) => order[slot] = action,
288 None => {
289 index.insert(canon, order.len());
290 order.push(action);
291 }
292 }
293 }
294
295 for action in order {
296 match action {
297 Action::Upsert(v) => plan.upserts.push(v),
298 Action::Delete(k) => plan.deletes.push(k),
299 }
300 }
301 plan
302}
303
304pub fn record_key(rec: &Value, key: &[String]) -> Option<KeyTuple> {
307 extract_key(rec, key).ok()
308}
309
310fn extract_key(rec: &Value, key: &[String]) -> Result<KeyTuple, String> {
313 let obj = rec
314 .as_object()
315 .ok_or_else(|| "record is not a JSON object".to_string())?;
316 let mut out = Vec::with_capacity(key.len());
317 for col in key {
318 match obj.get(col) {
319 None => return Err(format!("missing key column '{col}'")),
320 Some(Value::Null) => return Err(format!("null value for key column '{col}'")),
321 Some(v) => out.push((col.clone(), v.clone())),
322 }
323 }
324 Ok(KeyTuple(out))
325}
326
327fn is_delete_marked(rec: &Value, marker: Option<&DeleteMarker>) -> bool {
328 let Some(dm) = marker else { return false };
329 let Some(v) = rec.get(&dm.field) else {
330 return false;
331 };
332 let Some(s) = v.as_str() else { return false };
333 dm.values.iter().any(|m| m == s)
334}
335
336fn strip_marker(mut rec: Value, marker: Option<&DeleteMarker>) -> Value {
337 if let (Some(dm), Value::Object(map)) = (marker, &mut rec) {
338 map.remove(&dm.field);
339 }
340 rec
341}
342
343fn canonical(k: &KeyTuple) -> String {
348 let arr: Vec<String> = k.0.iter().map(|(_, v)| key_scalar_text(v)).collect();
349 serde_json::to_string(&arr).expect("a Vec<String> always serializes")
350}
351
352fn key_scalar_text(v: &Value) -> String {
353 match v {
354 Value::String(s) => s.clone(),
355 Value::Number(n) => match n.as_f64() {
356 Some(f)
357 if n.as_i64().is_none()
358 && n.as_u64().is_none()
359 && f.fract() == 0.0
360 && f.abs() < 9_007_199_254_740_992.0 =>
361 {
362 (f as i64).to_string()
363 }
364 _ => n.to_string(),
365 },
366 Value::Bool(b) => b.to_string(),
367 other => crate::util::canonical_json(other),
368 }
369}
370
371pub fn key_to_doc_id(k: &KeyTuple, separator: &str) -> String {
385 let _ = separator; if k.0.len() == 1 {
387 return match &k.0[0].1 {
388 Value::String(s) => s.clone(),
389 other => crate::util::canonical_json(other),
390 };
391 }
392 let values: Vec<Value> = k.0.iter().map(|(_, v)| v.clone()).collect();
393 crate::util::canonical_json(&Value::Array(values))
394}
395
396pub fn key_to_filter(k: &KeyTuple) -> Map<String, Value> {
398 k.0.iter().map(|(c, v)| (c.clone(), v.clone())).collect()
399}
400
401#[cfg(test)]
402mod tests {
403 use super::*;
404 use serde_json::json;
405
406 fn login_then_id() -> Value {
408 let mut m = serde_json::Map::new();
409 m.insert("login".into(), json!("o"));
410 m.insert("id".into(), json!(2));
411 Value::Object(m)
412 }
413
414 #[test]
415 fn object_key_values_render_with_sorted_keys() {
416 assert_eq!(key_scalar_text(&login_then_id()), r#"{"id":2,"login":"o"}"#);
417 let k = KeyTuple(vec![("k".into(), login_then_id())]);
418 assert_eq!(key_to_doc_id(&k, "_"), r#"{"id":2,"login":"o"}"#);
419 }
420
421 fn upsert_spec(keys: &[&str]) -> WriteSpec {
422 WriteSpec {
423 write_mode: WriteMode::Upsert,
424 key: keys.iter().map(|s| s.to_string()).collect(),
425 delete_marker: None,
426 rollback: None,
427 }
428 }
429
430 #[test]
431 fn upsert_extracts_key_and_keeps_row() {
432 let plan = plan_writes(&[json!({"id": 1, "name": "a"})], &upsert_spec(&["id"]));
433 assert_eq!(plan.upserts, vec![json!({"id": 1, "name": "a"})]);
434 assert!(plan.deletes.is_empty());
435 assert!(plan.failed.is_empty());
436 }
437
438 #[test]
439 fn key_to_doc_id_single_key_is_plain() {
440 let k = KeyTuple(vec![("id".into(), json!(7))]);
441 assert_eq!(key_to_doc_id(&k, "_"), "7");
442 let k = KeyTuple(vec![("name".into(), json!("alice"))]);
443 assert_eq!(key_to_doc_id(&k, "_"), "alice");
444 }
445
446 #[test]
447 fn key_to_doc_id_composite_is_injective() {
448 let k1 = KeyTuple(vec![("x".into(), json!("a_")), ("y".into(), json!("b"))]);
450 let k2 = KeyTuple(vec![("x".into(), json!("a")), ("y".into(), json!("_b"))]);
451 let id1 = key_to_doc_id(&k1, "_");
452 let id2 = key_to_doc_id(&k2, "_");
453 assert_ne!(id1, id2, "distinct composite keys must map to distinct ids");
454 let k3 = KeyTuple(vec![("x".into(), json!(1)), ("y".into(), json!("2"))]);
456 let k4 = KeyTuple(vec![("x".into(), json!("1")), ("y".into(), json!(2))]);
457 assert_ne!(key_to_doc_id(&k3, "_"), key_to_doc_id(&k4, "_"));
458 }
459
460 #[test]
461 fn missing_key_goes_to_failed_with_original_index() {
462 let plan = plan_writes(
463 &[json!({"id": 1}), json!({"name": "no-key"})],
464 &upsert_spec(&["id"]),
465 );
466 assert_eq!(plan.upserts.len(), 1);
467 assert_eq!(plan.failed.len(), 1);
468 assert_eq!(plan.failed[0].0, 1, "failed row keeps its page index");
469 }
470
471 #[test]
472 fn null_key_value_is_a_failure() {
473 let plan = plan_writes(&[json!({"id": null})], &upsert_spec(&["id"]));
474 assert!(plan.upserts.is_empty());
475 assert_eq!(plan.failed.len(), 1);
476 }
477
478 #[test]
479 fn delete_marker_routes_to_deletes_and_strips_marker() {
480 let spec = WriteSpec {
481 write_mode: WriteMode::Upsert,
482 key: vec!["id".into()],
483 delete_marker: Some(DeleteMarker {
484 field: "__op".into(),
485 values: vec!["d".into()],
486 }),
487 rollback: None,
488 };
489 let plan = plan_writes(
490 &[
491 json!({"id": 1, "name": "a", "__op": "u"}),
492 json!({"id": 2, "__op": "d"}),
493 ],
494 &spec,
495 );
496 assert_eq!(plan.upserts, vec![json!({"id": 1, "name": "a"})]);
497 assert_eq!(plan.deletes.len(), 1);
498 assert_eq!(plan.deletes[0].0, vec![("id".to_string(), json!(2))]);
499 }
500
501 #[test]
502 fn last_write_wins_dedup_keeps_final_upsert() {
503 let plan = plan_writes(
504 &[json!({"id": 1, "v": "old"}), json!({"id": 1, "v": "new"})],
505 &upsert_spec(&["id"]),
506 );
507 assert_eq!(plan.upserts, vec![json!({"id": 1, "v": "new"})]);
508 }
509
510 #[test]
511 fn mixed_json_types_of_one_key_dedup_as_one_key() {
512 let plan = plan_writes(
513 &[
514 json!({"id": 7, "v": "a"}),
515 json!({"id": 7.0, "v": "b"}),
516 json!({"id": "7", "v": "c"}),
517 json!({"id": 7.5, "v": "d"}),
518 ],
519 &upsert_spec(&["id"]),
520 );
521 assert_eq!(
522 plan.upserts,
523 vec![json!({"id": "7", "v": "c"}), json!({"id": 7.5, "v": "d"})]
524 );
525 }
526
527 #[test]
528 fn last_write_wins_delete_after_upsert_is_a_delete() {
529 let spec = WriteSpec {
530 write_mode: WriteMode::Upsert,
531 key: vec!["id".into()],
532 delete_marker: Some(DeleteMarker {
533 field: "__op".into(),
534 values: vec!["d".into()],
535 }),
536 rollback: None,
537 };
538 let plan = plan_writes(
539 &[json!({"id": 1, "__op": "u"}), json!({"id": 1, "__op": "d"})],
540 &spec,
541 );
542 assert!(plan.upserts.is_empty());
543 assert_eq!(plan.deletes.len(), 1);
544 }
545
546 #[test]
547 fn delete_mode_routes_every_row_to_deletes() {
548 let spec = WriteSpec {
549 write_mode: WriteMode::Delete,
550 key: vec!["id".into()],
551 delete_marker: None,
552 rollback: None,
553 };
554 let plan = plan_writes(&[json!({"id": 1}), json!({"id": 2})], &spec);
555 assert!(plan.upserts.is_empty());
556 assert_eq!(plan.deletes.len(), 2);
557 }
558
559 #[test]
560 fn composite_key_tuple_is_ordered() {
561 let plan = plan_writes(
562 &[json!({"a": 1, "b": 2, "v": 9})],
563 &upsert_spec(&["a", "b"]),
564 );
565 assert_eq!(plan.upserts.len(), 1);
566 let plan2 = plan_writes(
567 &[
568 json!({"a": 1, "b": 2, "v": "x"}),
569 json!({"a": 1, "b": 3, "v": "y"}),
570 ],
571 &upsert_spec(&["a", "b"]),
572 );
573 assert_eq!(plan2.upserts.len(), 2, "(1,2) and (1,3) are distinct keys");
574 }
575
576 #[test]
577 fn validate_rejects_upsert_without_key() {
578 let spec = WriteSpec {
579 write_mode: WriteMode::Upsert,
580 key: vec![],
581 delete_marker: None,
582 rollback: None,
583 };
584 assert!(spec.validate().is_err());
585 }
586
587 #[test]
588 fn validate_allows_append_without_key() {
589 assert!(WriteSpec::default().validate().is_ok());
590 }
591
592 #[test]
593 fn dedups_by_key_requires_keyed_upsert_or_delete() {
594 assert!(!WriteSpec::default().dedups_by_key());
595 let upsert = WriteSpec {
596 write_mode: WriteMode::Upsert,
597 key: vec!["id".into()],
598 delete_marker: None,
599 rollback: None,
600 };
601 assert!(upsert.dedups_by_key());
602 let delete = WriteSpec {
603 write_mode: WriteMode::Delete,
604 key: vec!["id".into()],
605 delete_marker: None,
606 rollback: None,
607 };
608 assert!(delete.dedups_by_key());
609 let keyless = WriteSpec {
611 write_mode: WriteMode::Upsert,
612 key: vec![],
613 delete_marker: None,
614 rollback: None,
615 };
616 assert!(!keyless.dedups_by_key());
617 }
618
619 #[test]
620 fn last_write_wins_upsert_after_delete_is_an_upsert() {
621 let spec = WriteSpec {
623 write_mode: WriteMode::Upsert,
624 key: vec!["id".into()],
625 delete_marker: Some(DeleteMarker {
626 field: "__op".into(),
627 values: vec!["d".into()],
628 }),
629 rollback: None,
630 };
631 let plan = plan_writes(
632 &[
633 json!({"id": 1, "__op": "d"}),
634 json!({"id": 1, "v": 9, "__op": "u"}),
635 ],
636 &spec,
637 );
638 assert!(plan.deletes.is_empty());
639 assert_eq!(plan.upserts, vec![json!({"id": 1, "v": 9})]);
640 }
641
642 #[test]
643 fn overwrite_mode_flags_and_needs_no_key() {
644 let spec = WriteSpec {
645 write_mode: WriteMode::Overwrite,
646 ..Default::default()
647 };
648 assert!(spec.is_overwrite());
649 assert!(!spec.dedups_by_key());
650 assert!(spec.validate().is_ok());
652 assert_eq!(WriteMode::Overwrite.as_str(), "overwrite");
653 assert!(!WriteSpec::default().is_overwrite());
655 assert!(!upsert_spec(&["id"]).is_overwrite());
656 }
657
658 #[test]
659 fn overwrite_scope_window_validates_and_renders() {
660 let scope = OverwriteScope::Window {
661 column: "posting_date".into(),
662 from: json!("2024-06-01"),
663 to: json!("2024-07-01"),
664 };
665 assert_eq!(scope.column(), "posting_date");
666 assert!(scope.validate().is_ok());
667 let whr = scope.render_where_literal("\"posting_date\"");
668 assert_eq!(
669 whr,
670 "\"posting_date\" >= '2024-06-01' AND \"posting_date\" < '2024-07-01'"
671 );
672
673 assert!(
675 OverwriteScope::Window {
676 column: " ".into(),
677 from: json!(1),
678 to: json!(2)
679 }
680 .validate()
681 .is_err()
682 );
683 assert!(
684 OverwriteScope::Window {
685 column: "c".into(),
686 from: json!(null),
687 to: json!(2)
688 }
689 .validate()
690 .is_err()
691 );
692 }
693
694 #[test]
695 fn scope_window_number_bounds() {
696 let scope = OverwriteScope::Window {
697 column: "seq".into(),
698 from: json!(100),
699 to: json!(200),
700 };
701 assert!(scope.validate().is_ok());
702 assert_eq!(
703 scope.render_where_literal("`seq`"),
704 "`seq` >= 100 AND `seq` < 200"
705 );
706 }
707
708 #[test]
709 fn sql_literal_renders_each_scalar_kind_by_the_ansi_rule() {
710 assert_eq!(sql_literal(&json!("plain")), "'plain'");
713 assert_eq!(sql_literal(&json!("O'Brien")), "'O''Brien'");
714 assert_eq!(sql_literal(&json!("a\\b")), "'a\\b'");
715 assert_eq!(sql_literal(&json!("'' ' ''")), "''''' '' '''''");
717 assert_eq!(sql_literal(&json!("")), "''");
719 assert_eq!(sql_literal(&json!(42)), "42");
721 assert_eq!(sql_literal(&json!(-1.5)), "-1.5");
722 assert_eq!(sql_literal(&json!(true)), "true");
723 assert_eq!(sql_literal(&json!(false)), "false");
724 assert_eq!(sql_literal(&Value::Null), "NULL");
726 assert_eq!(sql_literal(&json!({"a": 1})), "NULL");
727 assert_eq!(sql_literal(&json!([1, 2])), "NULL");
728 }
729
730 #[test]
731 fn sql_literal_neutralises_a_quote_breakout_attempt() {
732 let out = sql_literal(&json!("x' OR 1=1 --"));
734 assert_eq!(out, "'x'' OR 1=1 --'");
735 let interior = &out[1..out.len() - 1];
738 assert!(
739 interior.matches('\'').count().is_multiple_of(2),
740 "unpaired quote inside literal: {out}"
741 );
742 }
743
744 #[test]
745 fn scope_literal_escapes_quotes() {
746 let scope = OverwriteScope::Window {
748 column: "c".into(),
749 from: json!("x' OR '1'='1"),
750 to: json!("z"),
751 };
752 let whr = scope.render_where_literal("\"c\"");
753 assert!(whr.contains("'x'' OR ''1''=''1'"), "{whr}");
754 }
755
756 #[test]
757 fn render_where_with_uses_the_dialect_renderer() {
758 let scope = OverwriteScope::Window {
759 column: "c".into(),
760 from: json!("a"),
761 to: json!(5),
762 };
763 let whr = scope.render_where_with("`c`", |v| format!("<{v}>"));
764 assert_eq!(whr, "`c` >= <\"a\"> AND `c` < <5>");
765 }
766
767 #[test]
768 fn overwrite_deserializes_from_wire() {
769 let spec: WriteSpec = serde_json::from_value(json!({"write_mode": "overwrite"})).unwrap();
770 assert_eq!(spec.write_mode, WriteMode::Overwrite);
771 assert!(spec.is_overwrite());
772 }
773
774 #[test]
775 fn empty_page_produces_empty_plan() {
776 let plan = plan_writes(&[], &upsert_spec(&["id"]));
777 assert!(plan.upserts.is_empty());
778 assert!(plan.deletes.is_empty());
779 assert!(plan.failed.is_empty());
780 }
781
782 #[test]
783 fn delete_mode_dedups_repeated_key() {
784 let spec = WriteSpec {
786 write_mode: WriteMode::Delete,
787 key: vec!["id".into()],
788 delete_marker: None,
789 rollback: None,
790 };
791 let plan = plan_writes(&[json!({"id": 1}), json!({"id": 1})], &spec);
792 assert_eq!(plan.deletes.len(), 1);
793 }
794}