1use crate::error::{CliError, CliResult};
30use chrono::{DateTime, Duration, FixedOffset, TimeZone, Utc};
31
32pub const MAX_UNITS: usize = 10_000;
35pub const WARN_UNITS: usize = 1_000;
37
38#[derive(Debug, Clone, PartialEq, Eq)]
42pub struct TimeChunk {
43 pub id: String,
46 pub start: DateTime<FixedOffset>,
49 pub end: DateTime<FixedOffset>,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub enum WindowStep {
63 Absolute(Duration),
65 Days(i64),
67 Weeks(i64),
69}
70
71impl std::fmt::Display for WindowStep {
72 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81 match self {
82 Self::Absolute(d) => write!(f, "{}", d.num_seconds()),
83 Self::Days(n) => write!(f, "{n}d"),
84 Self::Weeks(n) => write!(f, "{n}w"),
85 }
86 }
87}
88
89impl WindowStep {
90 fn nominal(self) -> Duration {
92 match self {
93 Self::Absolute(d) => d,
94 Self::Days(n) => Duration::days(n),
95 Self::Weeks(n) => Duration::weeks(n),
96 }
97 }
98}
99
100pub fn parse_window(s: &str) -> CliResult<WindowStep> {
106 let s = s.trim();
107 let err = || {
108 CliError::Config(format!(
109 "'{s}' is not a valid window — use e.g. 45s, 30m, 6h, 1d, 1w"
110 ))
111 };
112 let (num, unit) = match s.chars().last() {
113 Some(c) if c.is_ascii_digit() => (s, "s"),
114 Some(c) => (&s[..s.len() - c.len_utf8()], &s[s.len() - c.len_utf8()..]),
115 None => return Err(err()),
116 };
117 let n: i64 = num.parse().map_err(|_| err())?;
118 if n <= 0 {
119 return Err(CliError::Config(format!(
120 "window '{s}' must be a positive duration"
121 )));
122 }
123 let step = match unit {
124 "s" => WindowStep::Absolute(Duration::seconds(n)),
125 "m" => WindowStep::Absolute(Duration::minutes(n)),
126 "h" => WindowStep::Absolute(Duration::hours(n)),
127 "d" => WindowStep::Days(n),
128 "w" => WindowStep::Weeks(n),
129 _ => return Err(err()),
130 };
131 Ok(step)
132}
133
134fn advance_calendar(cursor: DateTime<Utc>, tz: chrono_tz::Tz, days: i64) -> Option<DateTime<Utc>> {
144 let naive = cursor
145 .with_timezone(&tz)
146 .naive_local()
147 .checked_add_signed(Duration::days(days))?;
148 for extra_hours in 0..=3 {
149 let candidate = naive.checked_add_signed(Duration::hours(extra_hours))?;
150 if let Some(local) = tz.from_local_datetime(&candidate).earliest() {
151 return Some(local.with_timezone(&Utc));
152 }
153 }
154 None
155}
156
157pub fn parse_boundary(s: &str, tz: chrono_tz::Tz) -> CliResult<DateTime<FixedOffset>> {
161 if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
162 return Ok(dt.with_timezone(&tz).fixed_offset());
163 }
164 if let Ok(date) = s.parse::<chrono::NaiveDate>() {
165 let midnight = date
166 .and_hms_opt(0, 0, 0)
167 .ok_or_else(|| CliError::Config(format!("'{s}' has no valid midnight in {tz}")))?;
168 let local = tz
169 .from_local_datetime(&midnight)
170 .earliest()
171 .ok_or_else(|| {
172 CliError::Config(format!("'{s}' midnight does not exist in {tz} (DST gap)"))
173 })?;
174 return Ok(local.fixed_offset());
175 }
176 Err(CliError::Config(format!(
177 "'{s}' is not RFC3339 (2026-06-01T00:00:00Z) or a date (2026-06-01)"
178 )))
179}
180
181pub fn plan_windows(
187 from: DateTime<FixedOffset>,
188 to: DateTime<FixedOffset>,
189 window: Option<WindowStep>,
190 tz: chrono_tz::Tz,
191) -> CliResult<Vec<TimeChunk>> {
192 if from >= to {
193 return Err(CliError::Config(format!(
194 "--from ({from}) must be before --to ({to})"
195 )));
196 }
197 let mut units = Vec::new();
198 let mut cursor = from.with_timezone(&Utc);
199 let end = to.with_timezone(&Utc);
200 let step = window.unwrap_or(WindowStep::Absolute(end - cursor));
201 while cursor < end {
202 if units.len() >= MAX_UNITS {
203 return Err(CliError::Config(format!(
204 "the range would produce more than {MAX_UNITS} units with this --window — \
205 use a larger window"
206 )));
207 }
208 let next = match step {
213 WindowStep::Absolute(d) => cursor + d,
214 WindowStep::Days(n) => {
215 advance_calendar(cursor, tz, n).unwrap_or(cursor + step.nominal())
216 }
217 WindowStep::Weeks(n) => {
218 advance_calendar(cursor, tz, n * 7).unwrap_or(cursor + step.nominal())
219 }
220 };
221 let next = if next > cursor {
222 next
223 } else {
224 cursor + step.nominal()
225 };
226 let unit_end = next.min(end);
227 units.push(TimeChunk {
228 id: cursor.format("%Y%m%dT%H%M%SZ").to_string(),
229 start: cursor.with_timezone(&tz).fixed_offset(),
230 end: unit_end.with_timezone(&tz).fixed_offset(),
231 });
232 cursor = unit_end;
233 }
234 Ok(units)
235}
236
237#[derive(
242 Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
243)]
244#[serde(rename_all = "snake_case")]
245pub enum Bounds {
246 Inclusive,
249 HalfOpen,
252}
253
254#[derive(Debug, Clone, PartialEq, Eq)]
256pub struct IntChunk {
257 pub id: String,
260 pub start: i64,
262 pub end: i64,
264 pub is_last: bool,
268}
269
270pub fn plan_int_chunks(
276 from: i64,
277 to: i64,
278 chunk_size: u64,
279 bounds: Bounds,
280) -> CliResult<Vec<IntChunk>> {
281 if chunk_size == 0 {
282 return Err(CliError::Config(
283 "partition.chunk_size must be greater than 0".into(),
284 ));
285 }
286 let span: i128 = match bounds {
288 Bounds::Inclusive => to as i128 - from as i128 + 1,
289 Bounds::HalfOpen => to as i128 - from as i128,
290 };
291 if span <= 0 {
292 return Err(CliError::Config(format!(
293 "partition range is empty: from ({from}) must be {} to ({to})",
294 match bounds {
295 Bounds::Inclusive => "less than or equal to",
296 Bounds::HalfOpen => "less than",
297 }
298 )));
299 }
300 let size = chunk_size as u128;
303 let count = (span as u128).div_ceil(size) as i128;
304 if count > MAX_UNITS as i128 {
305 return Err(CliError::Config(format!(
306 "the range would produce {count} chunks with chunk_size {chunk_size} \
307 (max {MAX_UNITS}) — use a larger chunk_size"
308 )));
309 }
310 let width = (count.max(1) - 1).to_string().len();
311
312 let mut out = Vec::with_capacity(count as usize);
313 let mut cursor = from as i128;
314 for i in 0..count {
315 let next = cursor + size as i128;
316 let is_last = i == count - 1;
317 let raw_end = match bounds {
320 Bounds::Inclusive => (next - 1).min(to as i128),
321 Bounds::HalfOpen => next.min(to as i128),
322 };
323 out.push(IntChunk {
324 id: format!("{:0width$}", i, width = width),
325 start: cursor as i64,
326 end: raw_end as i64,
327 is_last,
328 });
329 cursor = next;
330 }
331 Ok(out)
332}
333
334#[derive(Debug, Clone, PartialEq, Eq)]
338pub struct OffsetChunk {
339 pub id: String,
340 pub offset: u64,
341 pub limit: u64,
342}
343
344pub fn plan_offset_chunks(total: u64, chunk_size: u64) -> CliResult<Vec<OffsetChunk>> {
350 if chunk_size == 0 {
351 return Err(CliError::Config(
352 "partition.chunk_size must be greater than 0".into(),
353 ));
354 }
355 if total == 0 {
356 return Ok(Vec::new());
357 }
358 let count = total.div_ceil(chunk_size);
359 if count > MAX_UNITS as u64 {
360 return Err(CliError::Config(format!(
361 "a total of {total} would produce {count} chunks with chunk_size {chunk_size} \
362 (max {MAX_UNITS}) — use a larger chunk_size"
363 )));
364 }
365 let width = (count - 1).to_string().len();
366 Ok((0..count)
367 .map(|i| OffsetChunk {
368 id: format!("{:0width$}", i, width = width),
369 offset: i * chunk_size,
370 limit: chunk_size.min(total - i * chunk_size),
371 })
372 .collect())
373}
374
375#[cfg(test)]
376mod tests {
377 use super::*;
378 use chrono::Timelike;
379
380 fn tz(name: &str) -> chrono_tz::Tz {
384 name.parse().unwrap()
385 }
386
387 #[test]
388 fn window_durations_parse() {
389 assert_eq!(
391 parse_window("45s").unwrap(),
392 WindowStep::Absolute(Duration::seconds(45))
393 );
394 assert_eq!(
395 parse_window("30m").unwrap(),
396 WindowStep::Absolute(Duration::minutes(30))
397 );
398 assert_eq!(
399 parse_window("6h").unwrap(),
400 WindowStep::Absolute(Duration::hours(6))
401 );
402 assert_eq!(
403 parse_window("3600").unwrap(),
404 WindowStep::Absolute(Duration::seconds(3600))
405 );
406 assert_eq!(parse_window("1d").unwrap(), WindowStep::Days(1));
408 assert_eq!(parse_window("2w").unwrap(), WindowStep::Weeks(2));
409 assert!(parse_window("0d").is_err());
410 assert!(parse_window("-1h").is_err());
411 assert!(parse_window("soon").is_err());
412 assert!(parse_window("1y").is_err());
413 }
414
415 #[test]
416 fn boundaries_parse_rfc3339_and_dates() {
417 let utc = tz("UTC");
418 let dt = parse_boundary("2026-06-01T12:30:00Z", utc).unwrap();
419 assert_eq!(dt.to_rfc3339(), "2026-06-01T12:30:00+00:00");
420 let ny = tz("America/New_York");
422 let dt = parse_boundary("2026-06-01", ny).unwrap();
423 assert_eq!(dt.to_rfc3339(), "2026-06-01T00:00:00-04:00");
424 assert!(parse_boundary("yesterday", utc).is_err());
425 }
426
427 #[test]
428 fn thirty_one_days_one_day_window_is_31_units() {
429 let utc = tz("UTC");
432 let from = parse_boundary("2026-06-01", utc).unwrap();
433 let to = parse_boundary("2026-07-02", utc).unwrap();
434 let units = plan_windows(from, to, Some(WindowStep::Days(1)), utc).unwrap();
435 assert_eq!(units.len(), 31);
436 assert_eq!(units[0].id, "20260601T000000Z");
437 assert_eq!(units[0].start.to_rfc3339(), "2026-06-01T00:00:00+00:00");
438 assert_eq!(units[0].end.to_rfc3339(), "2026-06-02T00:00:00+00:00");
439 for w in units.windows(2) {
441 assert_eq!(w[0].end, w[1].start);
442 }
443 assert_eq!(units[30].end.to_rfc3339(), "2026-07-02T00:00:00+00:00");
444 }
445
446 #[test]
447 fn last_window_truncates_at_to() {
448 let utc = tz("UTC");
449 let from = parse_boundary("2026-06-01T00:00:00Z", utc).unwrap();
450 let to = parse_boundary("2026-06-01T05:30:00Z", utc).unwrap();
451 let units = plan_windows(
452 from,
453 to,
454 Some(WindowStep::Absolute(Duration::hours(2))),
455 utc,
456 )
457 .unwrap();
458 assert_eq!(units.len(), 3);
459 assert_eq!(units[2].start.to_rfc3339(), "2026-06-01T04:00:00+00:00");
460 assert_eq!(units[2].end.to_rfc3339(), "2026-06-01T05:30:00+00:00");
461 }
462
463 #[test]
464 fn no_window_is_a_single_unit() {
465 let utc = tz("UTC");
466 let from = parse_boundary("2026-06-01", utc).unwrap();
467 let to = parse_boundary("2026-07-01", utc).unwrap();
468 let units = plan_windows(from, to, None, utc).unwrap();
469 assert_eq!(units.len(), 1);
470 assert_eq!(units[0].start, from);
471 assert_eq!(units[0].end, to);
472 }
473
474 #[test]
480 fn calendar_day_windows_stay_on_local_midnight_across_dst() {
481 let ny = tz("America/New_York");
482 let from = parse_boundary("2026-03-07", ny).unwrap();
483 let to = parse_boundary("2026-03-11", ny).unwrap();
484 let units = plan_windows(from, to, Some(WindowStep::Days(1)), ny).unwrap();
485
486 assert_eq!(units.len(), 4, "four calendar days");
487 for u in &units {
488 assert_eq!(
489 (u.start.hour(), u.start.minute()),
490 (0, 0),
491 "unit {} must start at local midnight, got {}",
492 u.id,
493 u.start
494 );
495 }
496 for w in units.windows(2) {
498 assert_eq!(w[0].end, w[1].start, "no gap/overlap");
499 }
500 let dates: Vec<String> = units
501 .iter()
502 .map(|u| u.start.format("%Y-%m-%d").to_string())
503 .collect();
504 assert_eq!(
505 dates,
506 ["2026-03-07", "2026-03-08", "2026-03-09", "2026-03-10"]
507 );
508 let spring_forward = &units[1];
510 assert_eq!(
511 (spring_forward.end - spring_forward.start).num_hours(),
512 23,
513 "2026-03-08 loses an hour"
514 );
515 }
516
517 #[test]
519 fn calendar_day_windows_handle_fall_back() {
520 let ny = tz("America/New_York");
521 let from = parse_boundary("2026-10-31", ny).unwrap();
522 let to = parse_boundary("2026-11-03", ny).unwrap();
523 let units = plan_windows(from, to, Some(WindowStep::Days(1)), ny).unwrap();
524 for u in &units {
525 assert_eq!((u.start.hour(), u.start.minute()), (0, 0), "{}", u.id);
526 }
527 let long_day = units
529 .iter()
530 .find(|u| u.start.format("%Y-%m-%d").to_string() == "2026-11-01")
531 .expect("the fall-back day is planned");
532 assert_eq!((long_day.end - long_day.start).num_hours(), 25);
533 }
534
535 #[test]
538 fn calendar_and_absolute_windows_differ_across_dst() {
539 let ny = tz("America/New_York");
540 let from = parse_boundary("2026-03-07", ny).unwrap();
541 let to = parse_boundary("2026-03-10", ny).unwrap();
542 let cal = plan_windows(from, to, Some(parse_window("1d").unwrap()), ny).unwrap();
543 let abs = plan_windows(from, to, Some(parse_window("24h").unwrap()), ny).unwrap();
544 assert_eq!(cal[2].start.hour(), 0, "calendar stays on midnight");
545 assert_eq!(abs[2].start.hour(), 1, "absolute drifts by the DST delta");
546 assert_ne!(cal[2].start, abs[2].start);
547 }
548
549 #[test]
552 fn window_descriptor_is_stable_for_absolute_and_distinct_for_calendar() {
553 assert_eq!(
554 WindowStep::Absolute(Duration::hours(6)).to_string(),
555 "21600"
556 );
557 assert_eq!(WindowStep::Absolute(Duration::days(1)).to_string(), "86400");
558 assert_eq!(WindowStep::Days(1).to_string(), "1d");
559 assert_eq!(WindowStep::Weeks(2).to_string(), "2w");
560 }
561
562 #[test]
563 fn dst_transition_produces_no_gap_or_overlap() {
564 let ny = tz("America/New_York");
567 let from = parse_boundary("2026-03-07", ny).unwrap();
568 let to = parse_boundary("2026-03-10T00:00:00-04:00", ny).unwrap();
569 let units =
570 plan_windows(from, to, Some(WindowStep::Absolute(Duration::days(1))), ny).unwrap();
571 for w in units.windows(2) {
572 assert_eq!(w[0].end, w[1].start, "no gap/overlap across DST");
573 }
574 assert!(units[0].start.to_rfc3339().ends_with("-05:00"));
576 assert!(units.last().unwrap().end.to_rfc3339().ends_with("-04:00"));
577 }
578
579 #[test]
580 fn rejects_inverted_range_and_unit_explosion() {
581 let utc = tz("UTC");
582 let from = parse_boundary("2026-06-02", utc).unwrap();
583 let to = parse_boundary("2026-06-01", utc).unwrap();
584 assert!(plan_windows(from, to, None, utc).is_err());
585
586 let from = parse_boundary("2020-01-01", utc).unwrap();
587 let to = parse_boundary("2026-01-01", utc).unwrap();
588 let err = plan_windows(
589 from,
590 to,
591 Some(WindowStep::Absolute(Duration::minutes(1))),
592 utc,
593 )
594 .unwrap_err();
595 assert!(err.to_string().contains("larger window"), "{err}");
596 }
597
598 fn covered(chunks: &[IntChunk], bounds: Bounds) -> Vec<i64> {
603 let mut seen = Vec::new();
604 for c in chunks {
605 let last = match bounds {
606 Bounds::Inclusive => c.end,
607 Bounds::HalfOpen => c.end - 1,
608 };
609 for v in c.start..=last {
610 seen.push(v);
611 }
612 }
613 seen
614 }
615
616 #[test]
617 fn inclusive_chunks_tile_the_range_exactly_once() {
618 let chunks = plan_int_chunks(0, 24, 10, Bounds::Inclusive).unwrap();
619 assert_eq!(chunks.len(), 3);
620 assert_eq!((chunks[0].start, chunks[0].end), (0, 9));
621 assert_eq!((chunks[1].start, chunks[1].end), (10, 19));
622 assert_eq!((chunks[2].start, chunks[2].end), (20, 24), "last truncated");
623 assert_eq!(
624 covered(&chunks, Bounds::Inclusive),
625 (0..=24).collect::<Vec<_>>()
626 );
627 }
628
629 #[test]
630 fn half_open_chunks_tile_the_range_exactly_once() {
631 let chunks = plan_int_chunks(0, 25, 10, Bounds::HalfOpen).unwrap();
632 assert_eq!(chunks.len(), 3);
633 assert_eq!((chunks[0].start, chunks[0].end), (0, 10));
634 assert_eq!((chunks[1].start, chunks[1].end), (10, 20));
635 assert_eq!((chunks[2].start, chunks[2].end), (20, 25));
636 assert_eq!(
637 covered(&chunks, Bounds::HalfOpen),
638 (0..25).collect::<Vec<_>>()
639 );
640 }
641
642 #[test]
643 fn the_two_bounds_differ_by_exactly_one_at_every_boundary() {
644 let inc = plan_int_chunks(0, 19, 10, Bounds::Inclusive).unwrap();
647 let half = plan_int_chunks(0, 20, 10, Bounds::HalfOpen).unwrap();
648 assert_eq!(inc[0].end, 9);
649 assert_eq!(half[0].end, 10);
650 assert_eq!(inc[0].end + 1, half[0].end);
651 }
652
653 #[test]
654 fn tiles_exactly_once_across_many_sizes_and_ranges() {
655 for from in [-7i64, 0, 5, 1000] {
656 for span in [1i64, 2, 7, 10, 33, 100] {
657 for size in [1u64, 2, 3, 10, 64] {
658 let to = from + span - 1;
659 let chunks = plan_int_chunks(from, to, size, Bounds::Inclusive).unwrap();
660 assert_eq!(
661 covered(&chunks, Bounds::Inclusive),
662 (from..=to).collect::<Vec<_>>(),
663 "inclusive from={from} span={span} size={size}"
664 );
665 let chunks =
666 plan_int_chunks(from, from + span, size, Bounds::HalfOpen).unwrap();
667 assert_eq!(
668 covered(&chunks, Bounds::HalfOpen),
669 (from..from + span).collect::<Vec<_>>(),
670 "half-open from={from} span={span} size={size}"
671 );
672 }
673 }
674 }
675 }
676
677 #[test]
678 fn a_single_value_range_is_one_chunk_inclusive_and_empty_half_open() {
679 let inc = plan_int_chunks(5, 5, 10, Bounds::Inclusive).unwrap();
680 assert_eq!(inc.len(), 1);
681 assert_eq!((inc[0].start, inc[0].end), (5, 5));
682 assert!(plan_int_chunks(5, 5, 10, Bounds::HalfOpen).is_err());
685 }
686
687 #[test]
688 fn only_the_final_chunk_is_marked_last() {
689 let chunks = plan_int_chunks(0, 29, 10, Bounds::Inclusive).unwrap();
690 assert_eq!(
691 chunks.iter().filter(|c| c.is_last).count(),
692 1,
693 "exactly one chunk carries the open-ended tail flag"
694 );
695 assert!(chunks.last().unwrap().is_last);
696 }
697
698 #[test]
699 fn ids_are_zero_padded_so_they_sort_in_plan_order() {
700 let chunks = plan_int_chunks(0, 99, 1, Bounds::Inclusive).unwrap();
701 let mut ids: Vec<&str> = chunks.iter().map(|c| c.id.as_str()).collect();
702 let planned = ids.clone();
703 ids.sort_unstable();
704 assert_eq!(ids, planned, "lexicographic order must match plan order");
705 }
706
707 #[test]
708 fn rejects_inverted_and_empty_ranges() {
709 assert!(plan_int_chunks(10, 5, 10, Bounds::Inclusive).is_err());
710 assert!(plan_int_chunks(10, 10, 10, Bounds::HalfOpen).is_err());
711 }
712
713 #[test]
714 fn rejects_zero_chunk_size() {
715 let err = plan_int_chunks(0, 10, 0, Bounds::Inclusive).unwrap_err();
718 assert!(err.to_string().contains("greater than 0"), "{err}");
719 }
720
721 #[test]
722 fn rejects_a_chunk_explosion() {
723 let err = plan_int_chunks(0, 10_000_000, 1, Bounds::Inclusive).unwrap_err();
724 let msg = err.to_string();
725 assert!(msg.contains("larger chunk_size"), "{msg}");
726 assert!(msg.contains(&MAX_UNITS.to_string()), "names the cap: {msg}");
727 }
728
729 #[test]
730 fn does_not_overflow_near_i64_bounds() {
731 let chunks = plan_int_chunks(i64::MAX - 5, i64::MAX, 2, Bounds::Inclusive).unwrap();
732 assert_eq!(covered(&chunks, Bounds::Inclusive).len(), 6);
733 let chunks = plan_int_chunks(i64::MIN, i64::MIN + 5, 2, Bounds::Inclusive).unwrap();
734 assert_eq!(covered(&chunks, Bounds::Inclusive).len(), 6);
735 }
736
737 #[test]
740 fn offset_chunks_cover_the_total_without_overrunning_it() {
741 let chunks = plan_offset_chunks(25, 10).unwrap();
742 assert_eq!(chunks.len(), 3);
743 assert_eq!((chunks[0].offset, chunks[0].limit), (0, 10));
744 assert_eq!((chunks[1].offset, chunks[1].limit), (10, 10));
745 assert_eq!(
746 (chunks[2].offset, chunks[2].limit),
747 (20, 5),
748 "final limit is trimmed to the remainder"
749 );
750 assert_eq!(chunks.iter().map(|c| c.limit).sum::<u64>(), 25);
751 }
752
753 #[test]
754 fn an_exact_multiple_produces_full_chunks() {
755 let chunks = plan_offset_chunks(30, 10).unwrap();
756 assert_eq!(chunks.len(), 3);
757 assert!(chunks.iter().all(|c| c.limit == 10));
758 }
759
760 #[test]
761 fn a_zero_total_plans_nothing() {
762 assert!(plan_offset_chunks(0, 10).unwrap().is_empty());
763 }
764
765 #[test]
766 fn offset_rejects_zero_chunk_size_and_explosions() {
767 assert!(plan_offset_chunks(10, 0).is_err());
768 assert!(plan_offset_chunks(10_000_000, 1).is_err());
769 }
770}