1use std::collections::HashMap;
40use std::path::Path;
41
42use crate::config::Config;
43use crate::error::Result;
44use crate::manifest::{MANIFEST_FILENAME, ManifestPart, PartStatus, RunManifest};
45use crate::plan::{
46 ExtractionStrategy, ReconcileReport, RepairAction, RepairOutcome, RepairPlan, RepairReport,
47 ResolvedRunPlan, build_plan,
48};
49use crate::source;
50use crate::state::StateStore;
51
52use super::RunSummary;
53use super::chunked::{ChunkSource, run_chunked_sequential};
54use super::reconcile_cmd;
55
56pub enum RepairOutputFormat {
58 Pretty,
60 Json(Option<String>),
62}
63
64pub enum RepairReportSource {
66 File(String),
68 Auto,
70}
71
72pub fn run_repair_command(
73 config_path: &str,
74 export_name: &str,
75 params: Option<&HashMap<String, String>>,
76 report_source: RepairReportSource,
77 execute: bool,
78 format: RepairOutputFormat,
79) -> Result<()> {
80 let config = Config::load_with_params(config_path, params)?;
81 let config_dir = Path::new(config_path)
82 .parent()
83 .unwrap_or_else(|| Path::new("."));
84
85 let export = config
86 .exports
87 .iter()
88 .find(|e| e.name == export_name)
89 .ok_or_else(|| anyhow::anyhow!("export '{}' not found in config", export_name))?;
90
91 let mut plan = build_plan(&config, export, config_dir, false, false, false, params)?;
92 if !matches!(plan.strategy, ExtractionStrategy::Chunked(_)) {
93 anyhow::bail!(
94 "repair: '{}' mode — only chunked exports are supported in v1 (Epic H)",
95 plan.strategy.mode_label()
96 );
97 }
98
99 let state_path = config_dir.join(".rivet_state.db");
100 let state = StateStore::open(state_path.to_str().unwrap_or(".rivet_state.db"))?;
101
102 let reconcile_report = load_or_build_reconcile(&plan, &state, report_source)?;
103 let repair_plan = RepairPlan::from_reconcile(&reconcile_report);
104
105 if !execute {
106 emit_plan(&repair_plan, &format)?;
107 return Ok(());
108 }
109
110 if repair_plan.is_empty() {
111 println!(
112 "repair: nothing to repair for '{}' (reconcile report is clean)",
113 export_name
114 );
115 return Ok(());
116 }
117
118 let report = execute_repair(&mut plan, &state, repair_plan)?;
119 emit_report(&report, &format)?;
120 Ok(())
121}
122
123fn load_or_build_reconcile(
124 plan: &ResolvedRunPlan,
125 state: &StateStore,
126 source: RepairReportSource,
127) -> Result<ReconcileReport> {
128 match source {
129 RepairReportSource::File(path) => {
130 let raw = std::fs::read_to_string(&path)
131 .map_err(|e| anyhow::anyhow!("cannot read reconcile report '{}': {}", path, e))?;
132 let r: ReconcileReport = serde_json::from_str(&raw)
133 .map_err(|e| anyhow::anyhow!("invalid reconcile report '{}': {}", path, e))?;
134 if r.export_name != plan.export_name {
135 anyhow::bail!(
136 "repair: reconcile report is for export '{}' but config targets '{}'",
137 r.export_name,
138 plan.export_name
139 );
140 }
141 Ok(r)
142 }
143 RepairReportSource::Auto => reconcile_cmd::reconcile_chunked_fresh(plan, state),
144 }
145}
146
147fn execute_repair(
148 plan: &mut ResolvedRunPlan,
149 state: &StateStore,
150 repair_plan: RepairPlan,
151) -> Result<RepairReport> {
152 let mut results: Vec<(RepairAction, RepairOutcome)> =
153 Vec::with_capacity(repair_plan.actions.len());
154
155 let run_id = state
161 .get_latest_chunk_run(&plan.export_name)?
162 .map(|(rid, _, _, _)| rid);
163
164 let mut src = source::create_source(&plan.source)?;
170 let mut summary = RunSummary::new(plan);
171
172 let dest = crate::destination::create_destination(&plan.destination)?;
180
181 let mut new_parts: Vec<ManifestPart> = Vec::new();
183
184 for a in &repair_plan.actions {
185 let (start, end) = match (a.start_key.parse::<i64>(), a.end_key.parse::<i64>()) {
186 (Ok(s), Ok(e)) => (s, e),
187 _ => {
188 results.push((
189 a.clone(),
190 RepairOutcome::Skipped {
191 reason: format!("unparseable chunk keys [{}..{}]", a.start_key, a.end_key),
192 },
193 ));
194 continue;
195 }
196 };
197
198 let rows_before = summary.total_rows;
199 let parts_before = summary.manifest_parts.len();
200 let outcome = run_chunked_sequential(
201 &mut *src,
202 plan,
203 &mut summary,
204 Some(state),
205 ChunkSource::Precomputed(vec![(start, end)]),
206 );
207 match outcome {
208 Ok(()) => {
209 let rows = summary.total_rows - rows_before;
210 let mut chunk_parts: Vec<ManifestPart> =
215 summary.manifest_parts[parts_before..].to_vec();
216
217 for p in &mut chunk_parts {
227 if let Some(renamed) = relabel_repair_chunk_index(&p.path, a.chunk_index) {
228 match dest.r#move(&p.path, &renamed) {
229 Ok(()) => p.path = renamed,
230 Err(e) => log::warn!(
231 "repair: chunk {} re-exported but could not rename \
232 '{}' → '{}' to carry the original chunk index \
233 (the file is durable under its chunk0 name): {:#}",
234 a.chunk_index,
235 p.path,
236 renamed,
237 e
238 ),
239 }
240 }
241 }
242
243 if let Some(rid) = &run_id {
251 let file_name = chunk_parts.last().map(|p| p.path.as_str());
252 if let Err(e) = state.complete_chunk_task(rid, a.chunk_index, rows, file_name) {
253 log::warn!(
257 "repair: chunk {} re-exported but chunk_task update failed — \
258 reconcile will still report the old mismatch: {:#}",
259 a.chunk_index,
260 e
261 );
262 }
263 } else {
264 log::warn!(
265 "repair: chunk {} re-exported but no chunk run is recorded for export \
266 '{}' — chunk_task could not be updated; reconcile will not converge",
267 a.chunk_index,
268 plan.export_name
269 );
270 }
271
272 new_parts.extend(chunk_parts);
273 results.push((a.clone(), RepairOutcome::Executed { rows_written: rows }));
274 }
275 Err(e) => {
276 let msg = crate::redact::redact_error(&e);
277 results.push((a.clone(), RepairOutcome::Failed { error: msg }));
278 }
279 }
280 }
281
282 if !new_parts.is_empty()
289 && let Err(e) = record_repair_parts_in_manifest(&plan.destination, &new_parts)
290 {
291 log::warn!(
292 "repair: re-exported parts were written but the destination manifest could not be \
293 updated (the files are durable; `rivet validate` may flag them as untracked): {:#}",
294 e
295 );
296 }
297
298 Ok(RepairReport::new(
299 repair_plan,
300 format!("repair-{}", chrono::Utc::now().format("%Y%m%dT%H%M%S")),
301 results,
302 ))
303}
304
305fn record_repair_parts_in_manifest(
315 destination: &crate::config::DestinationConfig,
316 new_parts: &[ManifestPart],
317) -> Result<()> {
318 let dest = crate::destination::create_destination(destination)?;
319
320 let raw = match dest.head(MANIFEST_FILENAME)? {
324 Some(_) => crate::pipeline::validate_manifest::read_capped(
325 &*dest,
326 MANIFEST_FILENAME,
327 crate::pipeline::validate_manifest::MANIFEST_MAX_BYTES,
328 )?,
329 None => anyhow::bail!(
330 "no manifest.json at the destination prefix — cannot record repair parts \
331 (was the original export finalized?)"
332 ),
333 };
334 let mut manifest: RunManifest = serde_json::from_slice(&raw)
335 .map_err(|e| anyhow::anyhow!("destination manifest.json is unparseable: {e}"))?;
336
337 let mut next_id = manifest.parts.iter().map(|p| p.part_id).max().unwrap_or(0) + 1;
339 for p in new_parts {
340 manifest.parts.push(ManifestPart {
341 part_id: next_id,
342 path: p.path.clone(),
343 rows: p.rows,
344 size_bytes: p.size_bytes,
345 content_fingerprint: p.content_fingerprint.clone(),
346 content_md5: p.content_md5.clone(),
347 status: PartStatus::Committed,
348 });
349 next_id += 1;
350 }
351
352 manifest.row_count = manifest.committed_rows();
355 manifest.part_count = manifest.committed_part_count() as u32;
356 manifest.finished_at = chrono::Utc::now().to_rfc3339();
357
358 crate::pipeline::manifest_writer::write_manifest(&*dest, &manifest)?;
373 Ok(())
374}
375
376fn relabel_repair_chunk_index(path: &str, original_chunk_index: i64) -> Option<String> {
391 if original_chunk_index == 0 {
392 return None;
393 }
394 let token = "_chunk0_";
395 let at = path.rfind(token)?;
396 Some(format!(
397 "{}_chunk{}_{}",
398 &path[..at],
399 original_chunk_index,
400 &path[at + token.len()..],
401 ))
402}
403
404fn emit_plan(plan: &RepairPlan, format: &RepairOutputFormat) -> Result<()> {
405 match format {
406 RepairOutputFormat::Pretty => print_plan_pretty(plan),
407 RepairOutputFormat::Json(None) => println!("{}", plan.to_json_pretty()?),
408 RepairOutputFormat::Json(Some(path)) => {
409 std::fs::write(path, plan.to_json_pretty()?)
410 .map_err(|e| anyhow::anyhow!("cannot write repair plan '{}': {}", path, e))?;
411 println!("Repair plan written to: {}", path);
412 }
413 }
414 Ok(())
415}
416
417fn emit_report(report: &RepairReport, format: &RepairOutputFormat) -> Result<()> {
418 match format {
419 RepairOutputFormat::Pretty => print_report_pretty(report),
420 RepairOutputFormat::Json(None) => println!("{}", report.to_json_pretty()?),
421 RepairOutputFormat::Json(Some(path)) => {
422 std::fs::write(path, report.to_json_pretty()?)
423 .map_err(|e| anyhow::anyhow!("cannot write repair report '{}': {}", path, e))?;
424 println!("Repair report written to: {}", path);
425 }
426 }
427 Ok(())
428}
429
430fn print_plan_pretty(plan: &RepairPlan) {
431 println!();
432 println!(" Export : {}", plan.export_name);
433 println!(" Reconcile run : {}", plan.reconcile_run_id);
434 println!(" Actions : {}", plan.actions.len());
435 for a in &plan.actions {
436 println!(
437 " • chunk {} [{}..{}] — {}",
438 a.chunk_index, a.start_key, a.end_key, a.reason
439 );
440 }
441 if !plan.skipped.is_empty() {
442 println!(" Skipped :");
443 for s in &plan.skipped {
444 println!(" • {s}");
445 }
446 }
447 if plan.is_empty() && plan.skipped.is_empty() {
448 println!(" (nothing to repair)");
449 }
450 println!();
451}
452
453fn print_report_pretty(report: &RepairReport) {
454 println!();
455 println!(" Export : {}", report.plan.export_name);
456 println!(" Repair run : {}", report.repair_run_id);
457 println!(
458 " Summary : planned {} · executed {} · skipped {} · failed {} · rows {}",
459 report.summary.planned,
460 report.summary.executed,
461 report.summary.skipped,
462 report.summary.failed,
463 report.summary.rows_written,
464 );
465 for (a, out) in &report.results {
466 let tag = match out {
467 RepairOutcome::Executed { rows_written } => format!("executed ({rows_written} rows)"),
468 RepairOutcome::Skipped { reason } => format!("skipped ({reason})"),
469 RepairOutcome::Failed { error } => format!("failed ({error})"),
470 };
471 println!(
472 " • chunk {} [{}..{}] — {tag}",
473 a.chunk_index, a.start_key, a.end_key
474 );
475 }
476 println!();
477}
478
479#[cfg(test)]
480mod tests {
481 use super::*;
482 use crate::plan::{PartitionKind, PartitionResult, ReconcileReport};
483
484 #[test]
485 fn plan_from_auto_would_derive_actions_from_reconcile() {
486 let partitions = vec![
488 PartitionResult::classify(
489 PartitionKind::Chunk,
490 "chunk 0 [1..100]".into(),
491 Some(100),
492 Some(100),
493 ),
494 PartitionResult::classify(
495 PartitionKind::Chunk,
496 "chunk 1 [101..200]".into(),
497 Some(100),
498 Some(90),
499 ),
500 ];
501 let r = ReconcileReport::new(
502 "orders".into(),
503 "rec-1".into(),
504 "chunked".into(),
505 partitions,
506 );
507 let plan = RepairPlan::from_reconcile(&r);
508 assert_eq!(plan.actions.len(), 1);
509 assert_eq!(plan.actions[0].chunk_index, 1);
510 }
511
512 #[test]
515 fn relabel_repair_chunk_index_rewrites_chunk0_to_original() {
516 let written = "orders_20260611_120000_chunk0_a1b2c3d4e5f6a7b8.parquet";
519 let renamed = relabel_repair_chunk_index(written, 2)
520 .expect("a non-zero chunk index must produce a renamed path");
521 assert_eq!(
522 renamed,
523 "orders_20260611_120000_chunk2_a1b2c3d4e5f6a7b8.parquet"
524 );
525 assert!(!renamed.contains("_chunk0_"), "no chunk0 token survives");
526 }
527
528 #[test]
529 fn relabel_repair_chunk_index_handles_rotated_part_suffix() {
530 let written = "orders_20260611_120000_chunk0_a1b2c3d4e5f6a7b8_p1.parquet";
533 let renamed = relabel_repair_chunk_index(written, 3).unwrap();
534 assert_eq!(
535 renamed,
536 "orders_20260611_120000_chunk3_a1b2c3d4e5f6a7b8_p1.parquet"
537 );
538 }
539
540 #[test]
541 fn relabel_repair_chunk_index_is_noop_for_chunk_zero() {
542 let written = "orders_20260611_120000_chunk0_a1b2c3d4e5f6a7b8.parquet";
544 assert!(relabel_repair_chunk_index(written, 0).is_none());
545 }
546
547 #[test]
548 fn relabel_repair_chunk_index_leaves_unexpected_shapes_untouched() {
549 assert!(relabel_repair_chunk_index("orders_no_chunk_token.parquet", 5).is_none());
552 }
553
554 #[test]
563 fn repair_updates_the_run_unique_manifest_copy_not_just_the_canonical() {
564 use crate::config::{DestinationConfig, DestinationType};
565 use crate::manifest::{
566 MANIFEST_VERSION, ManifestDestination, ManifestSource, ManifestStatus,
567 run_unique_manifest_name,
568 };
569
570 let dir = tempfile::tempdir().unwrap();
571 let dpath = dir.path().to_str().unwrap().to_string();
572 let run_id = "orders_20260722T120000.000";
573 let part = |id: u32, rows: i64| ManifestPart {
574 part_id: id,
575 path: format!("orders_{id}.parquet"),
576 rows,
577 size_bytes: 100,
578 content_fingerprint: "xxh3:0000000000000000".into(),
579 content_md5: String::new(),
580 status: PartStatus::Committed,
581 };
582 let manifest = RunManifest {
584 manifest_version: MANIFEST_VERSION,
585 run_id: run_id.into(),
586 export_name: "public.orders".into(),
587 mode: "chunked".into(),
588 started_at: "2026-07-22T12:00:00Z".into(),
589 finished_at: "2026-07-22T12:00:10Z".into(),
590 status: ManifestStatus::Success,
591 source: ManifestSource {
592 engine: "postgres".into(),
593 schema: Some("public".into()),
594 table: Some("orders".into()),
595 extraction: None,
596 },
597 destination: ManifestDestination {
598 kind: "local".into(),
599 uri: dpath.clone(),
600 },
601 format: "parquet".into(),
602 compression: "zstd".into(),
603 schema_fingerprint: "xxh3:0123456789abcdef".into(),
604 row_count: 20,
605 part_count: 2,
606 parts: vec![part(1, 10), part(2, 10)],
607 column_checksums: None,
608 checksum_key_column: None,
609 };
610 let bytes = serde_json::to_vec_pretty(&manifest).unwrap();
611 std::fs::write(dir.path().join(MANIFEST_FILENAME), &bytes).unwrap();
614 std::fs::write(dir.path().join(run_unique_manifest_name(run_id)), &bytes).unwrap();
615
616 let dest_cfg = DestinationConfig {
617 destination_type: DestinationType::Local,
618 path: Some(dpath),
619 ..Default::default()
620 };
621 record_repair_parts_in_manifest(&dest_cfg, &[part(3, 7)]).unwrap();
623
624 let read = |name: String| -> RunManifest {
625 serde_json::from_slice(&std::fs::read(dir.path().join(name)).unwrap()).unwrap()
626 };
627 let run_unique = read(run_unique_manifest_name(run_id));
628 assert_eq!(
629 run_unique.parts.len(),
630 3,
631 "the loader-authoritative run-unique copy must list the repair part"
632 );
633 assert_eq!(
634 run_unique.row_count, 27,
635 "run-unique row_count must reflect the recovered rows (20 + 7)"
636 );
637 assert_eq!(read(MANIFEST_FILENAME.to_string()).parts.len(), 3);
639 }
640}