1use std::time::{Duration, Instant};
22
23use serde::Serialize;
24
25use crate::theme::Theme;
26
27#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
29pub enum BatchState {
30 #[default]
32 Normal,
33 Complete,
35 Warning,
37 Error,
39}
40
41#[derive(Debug, Clone)]
68pub struct BatchOperationTracker {
69 operation_name: String,
71 total_batches: u64,
73 completed_batches: u64,
75 total_rows: u64,
77 processed_rows: u64,
79 error_count: u64,
81 error_threshold: u64,
83 started_at: Instant,
85 batch_times: Vec<Duration>,
87 batch_rows: Vec<u64>,
89 smoothing_window: usize,
91 last_batch_start: Instant,
93 state: BatchState,
95 theme: Option<Theme>,
97 width: Option<usize>,
99}
100
101impl BatchOperationTracker {
102 #[must_use]
109 pub fn new(operation_name: impl Into<String>, total_batches: u64, total_rows: u64) -> Self {
110 let now = Instant::now();
111 Self {
112 operation_name: operation_name.into(),
113 total_batches,
114 completed_batches: 0,
115 total_rows,
116 processed_rows: 0,
117 error_count: 0,
118 error_threshold: 10,
119 started_at: now,
120 batch_times: Vec::with_capacity(10),
121 batch_rows: Vec::with_capacity(10),
122 smoothing_window: 5,
123 last_batch_start: now,
124 state: BatchState::Normal,
125 theme: None,
126 width: None,
127 }
128 }
129
130 #[must_use]
132 pub fn theme(mut self, theme: Theme) -> Self {
133 self.theme = Some(theme);
134 self
135 }
136
137 #[must_use]
141 pub fn error_threshold(mut self, threshold: u64) -> Self {
142 self.error_threshold = threshold;
143 self
144 }
145
146 #[must_use]
148 pub fn width(mut self, width: usize) -> Self {
149 self.width = Some(width);
150 self
151 }
152
153 #[must_use]
155 pub fn smoothing_window(mut self, size: usize) -> Self {
156 self.smoothing_window = size.max(1);
157 self
158 }
159
160 pub fn complete_batch(&mut self, rows_in_batch: u64) {
165 let now = Instant::now();
166 let duration = now.duration_since(self.last_batch_start);
167
168 self.batch_times.push(duration);
170 self.batch_rows.push(rows_in_batch);
171
172 while self.batch_times.len() > self.smoothing_window {
174 self.batch_times.remove(0);
175 self.batch_rows.remove(0);
176 }
177
178 self.completed_batches += 1;
179 self.processed_rows += rows_in_batch;
180 self.last_batch_start = now;
181
182 self.update_state();
183 }
184
185 pub fn record_error(&mut self) {
187 self.error_count += 1;
188 self.update_state();
189 }
190
191 pub fn record_errors(&mut self, count: u64) {
193 self.error_count += count;
194 self.update_state();
195 }
196
197 #[must_use]
199 pub fn operation_name(&self) -> &str {
200 &self.operation_name
201 }
202
203 #[must_use]
205 pub fn completed_batches(&self) -> u64 {
206 self.completed_batches
207 }
208
209 #[must_use]
211 pub fn total_batches(&self) -> u64 {
212 self.total_batches
213 }
214
215 #[must_use]
217 pub fn processed_rows(&self) -> u64 {
218 self.processed_rows
219 }
220
221 #[must_use]
223 pub fn total_rows(&self) -> u64 {
224 self.total_rows
225 }
226
227 #[must_use]
229 pub fn error_count(&self) -> u64 {
230 self.error_count
231 }
232
233 #[must_use]
235 pub fn current_state(&self) -> BatchState {
236 self.state
237 }
238
239 #[must_use]
241 pub fn is_complete(&self) -> bool {
242 self.completed_batches >= self.total_batches
243 }
244
245 #[must_use]
247 pub fn batch_percentage(&self) -> f64 {
248 if self.total_batches == 0 {
249 return 100.0;
250 }
251 (self.completed_batches as f64 / self.total_batches as f64) * 100.0
252 }
253
254 #[must_use]
256 pub fn row_percentage(&self) -> f64 {
257 if self.total_rows == 0 {
258 return 100.0;
259 }
260 (self.processed_rows as f64 / self.total_rows as f64) * 100.0
261 }
262
263 #[must_use]
265 pub fn elapsed_secs(&self) -> f64 {
266 self.started_at.elapsed().as_secs_f64()
267 }
268
269 #[must_use]
273 pub fn throughput(&self) -> f64 {
274 if self.batch_times.is_empty() {
275 let elapsed = self.elapsed_secs();
277 if elapsed < 0.001 {
278 return 0.0;
279 }
280 return self.processed_rows as f64 / elapsed;
281 }
282
283 let total_duration: Duration = self.batch_times.iter().sum();
284 let total_rows: u64 = self.batch_rows.iter().sum();
285
286 let secs = total_duration.as_secs_f64();
287 if secs < 0.001 {
288 return 0.0;
289 }
290
291 total_rows as f64 / secs
292 }
293
294 #[must_use]
296 pub fn success_rate(&self) -> f64 {
297 let total = self.processed_rows + self.error_count;
298 if total == 0 {
299 return 100.0;
300 }
301 (self.processed_rows as f64 / total as f64) * 100.0
302 }
303
304 fn update_state(&mut self) {
306 if self.completed_batches >= self.total_batches {
307 self.state = BatchState::Complete;
308 } else if self.error_count > self.error_threshold {
309 self.state = BatchState::Error;
310 } else if self.error_count > 0 {
311 self.state = BatchState::Warning;
312 } else {
313 self.state = BatchState::Normal;
314 }
315 }
316
317 #[must_use]
321 pub fn render_plain(&self) -> String {
322 let pct = self.batch_percentage();
323 let rate = self.throughput();
324
325 let mut parts = vec![format!(
326 "{}: {:.0}% ({}/{} batches), {}/{} rows",
327 self.operation_name,
328 pct,
329 self.completed_batches,
330 self.total_batches,
331 self.processed_rows,
332 self.total_rows
333 )];
334
335 if self.processed_rows > 0 {
336 parts.push(format!("{rate:.0} rows/s"));
337 }
338
339 parts.push(format!("{} errors", self.error_count));
340
341 parts.join(", ")
342 }
343
344 #[must_use]
348 #[allow(clippy::cast_possible_truncation)] pub fn render_styled(&self) -> String {
350 let bar_width = self.width.unwrap_or(30);
351 let pct = self.batch_percentage();
352 let filled = ((pct / 100.0) * bar_width as f64).round() as usize;
353 let empty = bar_width.saturating_sub(filled);
354
355 let theme = self.theme.clone().unwrap_or_default();
356
357 let (bar_color, text_color) = match self.state {
358 BatchState::Normal => (theme.info.color_code(), theme.info.color_code()),
359 BatchState::Complete => (theme.success.color_code(), theme.success.color_code()),
360 BatchState::Warning => (theme.warning.color_code(), theme.warning.color_code()),
361 BatchState::Error => (theme.error.color_code(), theme.error.color_code()),
362 };
363 let reset = "\x1b[0m";
364
365 let bar = format!(
367 "{bar_color}[{filled}{empty}]{reset}",
368 filled = "=".repeat(filled.saturating_sub(1)) + if filled > 0 { ">" } else { "" },
369 empty = " ".repeat(empty),
370 );
371
372 let line1 = format!(
373 "{text_color}{}{reset} {bar} {pct:.0}% ({}/{} batches)",
374 self.operation_name, self.completed_batches, self.total_batches
375 );
376
377 let rate = self.throughput();
379 let error_str = if self.error_count == 0 {
380 format!(
381 "{}{} errors{reset}",
382 theme.success.color_code(),
383 self.error_count
384 )
385 } else if self.error_count > self.error_threshold {
386 format!(
387 "{}{} errors (threshold exceeded!){reset}",
388 theme.error.color_code(),
389 self.error_count
390 )
391 } else {
392 format!(
393 "{}{} errors{reset}",
394 theme.warning.color_code(),
395 self.error_count
396 )
397 };
398
399 let line2 = format!(
400 " Rows: {}/{} | Rate: {:.0} rows/s | {}",
401 self.processed_rows, self.total_rows, rate, error_str
402 );
403
404 format!("{line1}\n{line2}")
405 }
406
407 #[must_use]
411 pub fn render_summary(&self) -> String {
412 let elapsed = self.elapsed_secs();
413 let avg_rate = if elapsed > 0.001 {
414 self.processed_rows as f64 / elapsed
415 } else {
416 0.0
417 };
418
419 format!(
420 "Summary for '{}':\n\
421 - Total time: {}\n\
422 - Total rows: {}\n\
423 - Average rate: {:.0} rows/s\n\
424 - Errors: {}\n\
425 - Success rate: {:.1}%",
426 self.operation_name,
427 format_duration(elapsed),
428 self.processed_rows,
429 avg_rate,
430 self.error_count,
431 self.success_rate()
432 )
433 }
434
435 #[must_use]
437 pub fn to_json(&self) -> String {
438 #[derive(Serialize)]
439 struct BatchJson<'a> {
440 operation: &'a str,
441 completed_batches: u64,
442 total_batches: u64,
443 processed_rows: u64,
444 total_rows: u64,
445 batch_percentage: f64,
446 row_percentage: f64,
447 throughput: f64,
448 error_count: u64,
449 error_threshold: u64,
450 elapsed_secs: f64,
451 is_complete: bool,
452 success_rate: f64,
453 state: &'a str,
454 }
455
456 let state_str = match self.state {
457 BatchState::Normal => "normal",
458 BatchState::Complete => "complete",
459 BatchState::Warning => "warning",
460 BatchState::Error => "error",
461 };
462
463 let json = BatchJson {
464 operation: &self.operation_name,
465 completed_batches: self.completed_batches,
466 total_batches: self.total_batches,
467 processed_rows: self.processed_rows,
468 total_rows: self.total_rows,
469 batch_percentage: self.batch_percentage(),
470 row_percentage: self.row_percentage(),
471 throughput: self.throughput(),
472 error_count: self.error_count,
473 error_threshold: self.error_threshold,
474 elapsed_secs: self.elapsed_secs(),
475 is_complete: self.is_complete(),
476 success_rate: self.success_rate(),
477 state: state_str,
478 };
479
480 serde_json::to_string(&json).unwrap_or_else(|_| "{}".to_string())
481 }
482}
483
484fn format_duration(secs: f64) -> String {
486 if secs < 1.0 {
487 return format!("{:.0}ms", secs * 1000.0);
488 }
489 if secs < 60.0 {
490 return format!("{secs:.1}s");
491 }
492 if secs < 3600.0 {
493 let mins = (secs / 60.0).floor();
494 let remaining = secs % 60.0;
495 return format!("{mins:.0}m{remaining:.0}s");
496 }
497 let hours = (secs / 3600.0).floor();
498 let remaining_mins = ((secs % 3600.0) / 60.0).floor();
499 format!("{hours:.0}h{remaining_mins:.0}m")
500}
501
502#[cfg(test)]
503mod tests {
504 use super::*;
505
506 #[test]
507 fn test_batch_tracker_creation() {
508 let tracker = BatchOperationTracker::new("Test", 10, 1000);
509 assert_eq!(tracker.operation_name(), "Test");
510 assert_eq!(tracker.total_batches(), 10);
511 assert_eq!(tracker.total_rows(), 1000);
512 assert_eq!(tracker.completed_batches(), 0);
513 assert_eq!(tracker.processed_rows(), 0);
514 assert_eq!(tracker.error_count(), 0);
515 assert_eq!(tracker.current_state(), BatchState::Normal);
516 }
517
518 #[test]
519 fn test_batch_complete() {
520 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
521 assert_eq!(tracker.completed_batches(), 0);
522
523 tracker.complete_batch(100);
524 assert_eq!(tracker.completed_batches(), 1);
525
526 tracker.complete_batch(100);
527 assert_eq!(tracker.completed_batches(), 2);
528 }
529
530 #[test]
531 fn test_batch_rows_tracking() {
532 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
533
534 tracker.complete_batch(100);
535 assert_eq!(tracker.processed_rows(), 100);
536
537 tracker.complete_batch(150);
538 assert_eq!(tracker.processed_rows(), 250);
539
540 tracker.complete_batch(50);
541 assert_eq!(tracker.processed_rows(), 300);
542 }
543
544 #[test]
545 fn test_batch_rate_calculation() {
546 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
547
548 assert!(tracker.throughput() >= 0.0);
550
551 tracker.complete_batch(100);
553
554 assert!(tracker.throughput() >= 0.0);
556 }
557
558 #[test]
559 fn test_batch_error_recording() {
560 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
561 assert_eq!(tracker.error_count(), 0);
562
563 tracker.record_error();
564 assert_eq!(tracker.error_count(), 1);
565
566 tracker.record_errors(5);
567 assert_eq!(tracker.error_count(), 6);
568 }
569
570 #[test]
571 fn test_batch_error_threshold() {
572 let mut tracker = BatchOperationTracker::new("Test", 10, 1000).error_threshold(5);
573 tracker.complete_batch(100);
574
575 assert_eq!(tracker.current_state(), BatchState::Normal);
577
578 tracker.record_errors(3);
580 assert_eq!(tracker.current_state(), BatchState::Warning);
581
582 tracker.record_errors(5);
584 assert_eq!(tracker.current_state(), BatchState::Error);
585 }
586
587 #[test]
588 fn test_batch_render_plain() {
589 let mut tracker = BatchOperationTracker::new("Batch insert", 20, 10000);
590 tracker.complete_batch(500);
591
592 let plain = tracker.render_plain();
593 assert!(plain.contains("Batch insert:"));
594 assert!(plain.contains("5%"));
595 assert!(plain.contains("(1/20 batches)"));
596 assert!(plain.contains("500/10000 rows"));
597 assert!(plain.contains("0 errors"));
598 }
599
600 #[test]
601 fn test_batch_render_plain_with_errors() {
602 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
603 tracker.complete_batch(100);
604 tracker.record_errors(3);
605
606 let plain = tracker.render_plain();
607 assert!(plain.contains("3 errors"));
608 }
609
610 #[test]
611 fn test_batch_summary() {
612 let mut tracker = BatchOperationTracker::new("Migration", 5, 500);
613 tracker.complete_batch(100);
614 tracker.complete_batch(100);
615 tracker.complete_batch(100);
616 tracker.complete_batch(100);
617 tracker.complete_batch(100);
618
619 let summary = tracker.render_summary();
620 assert!(summary.contains("Migration"));
621 assert!(summary.contains("Total rows: 500"));
622 assert!(summary.contains("Errors: 0"));
623 assert!(summary.contains("Success rate:"));
624 }
625
626 #[test]
627 fn test_batch_single_batch() {
628 let mut tracker = BatchOperationTracker::new("Single", 1, 100);
629 tracker.complete_batch(100);
630
631 assert!(tracker.is_complete());
632 assert!((tracker.batch_percentage() - 100.0).abs() < f64::EPSILON);
633 assert_eq!(tracker.current_state(), BatchState::Complete);
634 }
635
636 #[test]
637 fn test_batch_many_batches() {
638 let mut tracker = BatchOperationTracker::new("Large", 100, 10000);
639
640 for _ in 0..100 {
641 tracker.complete_batch(100);
642 }
643
644 assert!(tracker.is_complete());
645 assert_eq!(tracker.processed_rows(), 10000);
646 assert_eq!(tracker.completed_batches(), 100);
647 }
648
649 #[test]
650 fn test_batch_percentage_calculation() {
651 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
652
653 assert!((tracker.batch_percentage() - 0.0).abs() < f64::EPSILON);
654
655 tracker.complete_batch(100);
656 assert!((tracker.batch_percentage() - 10.0).abs() < f64::EPSILON);
657
658 tracker.complete_batch(100);
659 tracker.complete_batch(100);
660 tracker.complete_batch(100);
661 tracker.complete_batch(100);
662 assert!((tracker.batch_percentage() - 50.0).abs() < f64::EPSILON);
663 }
664
665 #[test]
666 fn test_row_percentage_calculation() {
667 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
668
669 assert!((tracker.row_percentage() - 0.0).abs() < f64::EPSILON);
670
671 tracker.complete_batch(250);
672 assert!((tracker.row_percentage() - 25.0).abs() < f64::EPSILON);
673
674 tracker.complete_batch(250);
675 assert!((tracker.row_percentage() - 50.0).abs() < f64::EPSILON);
676 }
677
678 #[test]
679 fn test_batch_zero_total() {
680 let tracker = BatchOperationTracker::new("Test", 0, 0);
681 assert!((tracker.batch_percentage() - 100.0).abs() < f64::EPSILON);
682 assert!((tracker.row_percentage() - 100.0).abs() < f64::EPSILON);
683 }
684
685 #[test]
686 fn test_batch_is_complete() {
687 let mut tracker = BatchOperationTracker::new("Test", 3, 300);
688 assert!(!tracker.is_complete());
689
690 tracker.complete_batch(100);
691 assert!(!tracker.is_complete());
692
693 tracker.complete_batch(100);
694 assert!(!tracker.is_complete());
695
696 tracker.complete_batch(100);
697 assert!(tracker.is_complete());
698 }
699
700 #[test]
701 fn test_batch_success_rate() {
702 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
703 tracker.complete_batch(100);
704
705 assert!((tracker.success_rate() - 100.0).abs() < 0.1);
707
708 tracker.record_error();
710 assert!(tracker.success_rate() > 99.0 && tracker.success_rate() < 100.0);
712 }
713
714 #[test]
715 fn test_batch_success_rate_no_data() {
716 let tracker = BatchOperationTracker::new("Test", 10, 1000);
717 assert!((tracker.success_rate() - 100.0).abs() < f64::EPSILON);
719 }
720
721 #[test]
722 fn test_batch_json_output() {
723 let mut tracker = BatchOperationTracker::new("Test", 10, 1000);
724 tracker.complete_batch(100);
725 tracker.record_error();
726
727 let json = tracker.to_json();
728 assert!(json.contains("\"operation\":\"Test\""));
729 assert!(json.contains("\"completed_batches\":1"));
730 assert!(json.contains("\"total_batches\":10"));
731 assert!(json.contains("\"processed_rows\":100"));
732 assert!(json.contains("\"error_count\":1"));
733 assert!(json.contains("\"state\":\"warning\""));
734 }
735
736 #[test]
737 fn test_batch_json_complete() {
738 let mut tracker = BatchOperationTracker::new("Test", 1, 100);
739 tracker.complete_batch(100);
740
741 let json = tracker.to_json();
742 assert!(json.contains("\"is_complete\":true"));
743 assert!(json.contains("\"state\":\"complete\""));
744 }
745
746 #[test]
747 fn test_batch_styled_contains_progress_bar() {
748 let mut tracker = BatchOperationTracker::new("Test", 10, 1000).width(20);
749 tracker.complete_batch(500);
750
751 let styled = tracker.render_styled();
752 assert!(styled.contains('['));
753 assert!(styled.contains(']'));
754 assert!(styled.contains("Rows:"));
755 }
756
757 #[test]
758 fn test_batch_styled_error_warning() {
759 let mut tracker = BatchOperationTracker::new("Test", 10, 1000)
760 .error_threshold(5)
761 .width(20);
762 tracker.complete_batch(100);
763 tracker.record_errors(10);
764
765 let styled = tracker.render_styled();
766 assert!(styled.contains("threshold exceeded"));
767 }
768
769 #[test]
770 fn test_batch_builder_chain() {
771 let tracker = BatchOperationTracker::new("Test", 10, 1000)
772 .theme(Theme::default())
773 .width(40)
774 .error_threshold(20)
775 .smoothing_window(10);
776
777 assert_eq!(tracker.total_batches(), 10);
778 }
779
780 #[test]
781 fn test_format_duration_ms() {
782 let result = format_duration(0.5);
783 assert!(result.contains("ms"));
784 }
785
786 #[test]
787 fn test_format_duration_seconds() {
788 let result = format_duration(30.0);
789 assert!(result.contains('s'));
790 assert!(!result.contains('m'));
791 }
792
793 #[test]
794 fn test_format_duration_minutes() {
795 let result = format_duration(125.0);
796 assert!(result.contains('m'));
797 }
798
799 #[test]
800 fn test_format_duration_hours() {
801 let result = format_duration(7300.0);
802 assert!(result.contains('h'));
803 }
804}