use crate::error::{CliError, CliResult};
use chrono::{DateTime, Duration, FixedOffset, TimeZone, Utc};
pub const MAX_UNITS: usize = 10_000;
pub const WARN_UNITS: usize = 1_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TimeChunk {
pub id: String,
pub start: DateTime<FixedOffset>,
pub end: DateTime<FixedOffset>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WindowStep {
Absolute(Duration),
Days(i64),
Weeks(i64),
}
impl std::fmt::Display for WindowStep {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Absolute(d) => write!(f, "{}", d.num_seconds()),
Self::Days(n) => write!(f, "{n}d"),
Self::Weeks(n) => write!(f, "{n}w"),
}
}
}
impl WindowStep {
fn nominal(self) -> Duration {
match self {
Self::Absolute(d) => d,
Self::Days(n) => Duration::days(n),
Self::Weeks(n) => Duration::weeks(n),
}
}
}
pub fn parse_window(s: &str) -> CliResult<WindowStep> {
let s = s.trim();
let err = || {
CliError::Config(format!(
"'{s}' is not a valid window — use e.g. 45s, 30m, 6h, 1d, 1w"
))
};
let (num, unit) = match s.chars().last() {
Some(c) if c.is_ascii_digit() => (s, "s"),
Some(c) => (&s[..s.len() - c.len_utf8()], &s[s.len() - c.len_utf8()..]),
None => return Err(err()),
};
let n: i64 = num.parse().map_err(|_| err())?;
if n <= 0 {
return Err(CliError::Config(format!(
"window '{s}' must be a positive duration"
)));
}
let step = match unit {
"s" => WindowStep::Absolute(Duration::seconds(n)),
"m" => WindowStep::Absolute(Duration::minutes(n)),
"h" => WindowStep::Absolute(Duration::hours(n)),
"d" => WindowStep::Days(n),
"w" => WindowStep::Weeks(n),
_ => return Err(err()),
};
Ok(step)
}
fn advance_calendar(cursor: DateTime<Utc>, tz: chrono_tz::Tz, days: i64) -> Option<DateTime<Utc>> {
let naive = cursor
.with_timezone(&tz)
.naive_local()
.checked_add_signed(Duration::days(days))?;
for extra_hours in 0..=3 {
let candidate = naive.checked_add_signed(Duration::hours(extra_hours))?;
if let Some(local) = tz.from_local_datetime(&candidate).earliest() {
return Some(local.with_timezone(&Utc));
}
}
None
}
pub fn parse_boundary(s: &str, tz: chrono_tz::Tz) -> CliResult<DateTime<FixedOffset>> {
if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
return Ok(dt.with_timezone(&tz).fixed_offset());
}
if let Ok(date) = s.parse::<chrono::NaiveDate>() {
let midnight = date
.and_hms_opt(0, 0, 0)
.ok_or_else(|| CliError::Config(format!("'{s}' has no valid midnight in {tz}")))?;
let local = tz
.from_local_datetime(&midnight)
.earliest()
.ok_or_else(|| {
CliError::Config(format!("'{s}' midnight does not exist in {tz} (DST gap)"))
})?;
return Ok(local.fixed_offset());
}
Err(CliError::Config(format!(
"'{s}' is not RFC3339 (2026-06-01T00:00:00Z) or a date (2026-06-01)"
)))
}
pub fn plan_windows(
from: DateTime<FixedOffset>,
to: DateTime<FixedOffset>,
window: Option<WindowStep>,
tz: chrono_tz::Tz,
) -> CliResult<Vec<TimeChunk>> {
if from >= to {
return Err(CliError::Config(format!(
"--from ({from}) must be before --to ({to})"
)));
}
let mut units = Vec::new();
let mut cursor = from.with_timezone(&Utc);
let end = to.with_timezone(&Utc);
let step = window.unwrap_or(WindowStep::Absolute(end - cursor));
while cursor < end {
if units.len() >= MAX_UNITS {
return Err(CliError::Config(format!(
"the range would produce more than {MAX_UNITS} units with this --window — \
use a larger window"
)));
}
let next = match step {
WindowStep::Absolute(d) => cursor + d,
WindowStep::Days(n) => {
advance_calendar(cursor, tz, n).unwrap_or(cursor + step.nominal())
}
WindowStep::Weeks(n) => {
advance_calendar(cursor, tz, n * 7).unwrap_or(cursor + step.nominal())
}
};
let next = if next > cursor {
next
} else {
cursor + step.nominal()
};
let unit_end = next.min(end);
units.push(TimeChunk {
id: cursor.format("%Y%m%dT%H%M%SZ").to_string(),
start: cursor.with_timezone(&tz).fixed_offset(),
end: unit_end.with_timezone(&tz).fixed_offset(),
});
cursor = unit_end;
}
Ok(units)
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum Bounds {
Inclusive,
HalfOpen,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IntChunk {
pub id: String,
pub start: i64,
pub end: i64,
pub is_last: bool,
}
pub fn plan_int_chunks(
from: i64,
to: i64,
chunk_size: u64,
bounds: Bounds,
) -> CliResult<Vec<IntChunk>> {
if chunk_size == 0 {
return Err(CliError::Config(
"partition.chunk_size must be greater than 0".into(),
));
}
let span: i128 = match bounds {
Bounds::Inclusive => to as i128 - from as i128 + 1,
Bounds::HalfOpen => to as i128 - from as i128,
};
if span <= 0 {
return Err(CliError::Config(format!(
"partition range is empty: from ({from}) must be {} to ({to})",
match bounds {
Bounds::Inclusive => "less than or equal to",
Bounds::HalfOpen => "less than",
}
)));
}
let size = chunk_size as u128;
let count = (span as u128).div_ceil(size) as i128;
if count > MAX_UNITS as i128 {
return Err(CliError::Config(format!(
"the range would produce {count} chunks with chunk_size {chunk_size} \
(max {MAX_UNITS}) — use a larger chunk_size"
)));
}
let width = (count.max(1) - 1).to_string().len();
let mut out = Vec::with_capacity(count as usize);
let mut cursor = from as i128;
for i in 0..count {
let next = cursor + size as i128;
let is_last = i == count - 1;
let raw_end = match bounds {
Bounds::Inclusive => (next - 1).min(to as i128),
Bounds::HalfOpen => next.min(to as i128),
};
out.push(IntChunk {
id: format!("{:0width$}", i, width = width),
start: cursor as i64,
end: raw_end as i64,
is_last,
});
cursor = next;
}
Ok(out)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OffsetChunk {
pub id: String,
pub offset: u64,
pub limit: u64,
}
pub fn plan_offset_chunks(total: u64, chunk_size: u64) -> CliResult<Vec<OffsetChunk>> {
if chunk_size == 0 {
return Err(CliError::Config(
"partition.chunk_size must be greater than 0".into(),
));
}
if total == 0 {
return Ok(Vec::new());
}
let count = total.div_ceil(chunk_size);
if count > MAX_UNITS as u64 {
return Err(CliError::Config(format!(
"a total of {total} would produce {count} chunks with chunk_size {chunk_size} \
(max {MAX_UNITS}) — use a larger chunk_size"
)));
}
let width = (count - 1).to_string().len();
Ok((0..count)
.map(|i| OffsetChunk {
id: format!("{:0width$}", i, width = width),
offset: i * chunk_size,
limit: chunk_size.min(total - i * chunk_size),
})
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Timelike;
fn tz(name: &str) -> chrono_tz::Tz {
name.parse().unwrap()
}
#[test]
fn window_durations_parse() {
assert_eq!(
parse_window("45s").unwrap(),
WindowStep::Absolute(Duration::seconds(45))
);
assert_eq!(
parse_window("30m").unwrap(),
WindowStep::Absolute(Duration::minutes(30))
);
assert_eq!(
parse_window("6h").unwrap(),
WindowStep::Absolute(Duration::hours(6))
);
assert_eq!(
parse_window("3600").unwrap(),
WindowStep::Absolute(Duration::seconds(3600))
);
assert_eq!(parse_window("1d").unwrap(), WindowStep::Days(1));
assert_eq!(parse_window("2w").unwrap(), WindowStep::Weeks(2));
assert!(parse_window("0d").is_err());
assert!(parse_window("-1h").is_err());
assert!(parse_window("soon").is_err());
assert!(parse_window("1y").is_err());
}
#[test]
fn boundaries_parse_rfc3339_and_dates() {
let utc = tz("UTC");
let dt = parse_boundary("2026-06-01T12:30:00Z", utc).unwrap();
assert_eq!(dt.to_rfc3339(), "2026-06-01T12:30:00+00:00");
let ny = tz("America/New_York");
let dt = parse_boundary("2026-06-01", ny).unwrap();
assert_eq!(dt.to_rfc3339(), "2026-06-01T00:00:00-04:00");
assert!(parse_boundary("yesterday", utc).is_err());
}
#[test]
fn thirty_one_days_one_day_window_is_31_units() {
let utc = tz("UTC");
let from = parse_boundary("2026-06-01", utc).unwrap();
let to = parse_boundary("2026-07-02", utc).unwrap();
let units = plan_windows(from, to, Some(WindowStep::Days(1)), utc).unwrap();
assert_eq!(units.len(), 31);
assert_eq!(units[0].id, "20260601T000000Z");
assert_eq!(units[0].start.to_rfc3339(), "2026-06-01T00:00:00+00:00");
assert_eq!(units[0].end.to_rfc3339(), "2026-06-02T00:00:00+00:00");
for w in units.windows(2) {
assert_eq!(w[0].end, w[1].start);
}
assert_eq!(units[30].end.to_rfc3339(), "2026-07-02T00:00:00+00:00");
}
#[test]
fn last_window_truncates_at_to() {
let utc = tz("UTC");
let from = parse_boundary("2026-06-01T00:00:00Z", utc).unwrap();
let to = parse_boundary("2026-06-01T05:30:00Z", utc).unwrap();
let units = plan_windows(
from,
to,
Some(WindowStep::Absolute(Duration::hours(2))),
utc,
)
.unwrap();
assert_eq!(units.len(), 3);
assert_eq!(units[2].start.to_rfc3339(), "2026-06-01T04:00:00+00:00");
assert_eq!(units[2].end.to_rfc3339(), "2026-06-01T05:30:00+00:00");
}
#[test]
fn no_window_is_a_single_unit() {
let utc = tz("UTC");
let from = parse_boundary("2026-06-01", utc).unwrap();
let to = parse_boundary("2026-07-01", utc).unwrap();
let units = plan_windows(from, to, None, utc).unwrap();
assert_eq!(units.len(), 1);
assert_eq!(units[0].start, from);
assert_eq!(units[0].end, to);
}
#[test]
fn calendar_day_windows_stay_on_local_midnight_across_dst() {
let ny = tz("America/New_York");
let from = parse_boundary("2026-03-07", ny).unwrap();
let to = parse_boundary("2026-03-11", ny).unwrap();
let units = plan_windows(from, to, Some(WindowStep::Days(1)), ny).unwrap();
assert_eq!(units.len(), 4, "four calendar days");
for u in &units {
assert_eq!(
(u.start.hour(), u.start.minute()),
(0, 0),
"unit {} must start at local midnight, got {}",
u.id,
u.start
);
}
for w in units.windows(2) {
assert_eq!(w[0].end, w[1].start, "no gap/overlap");
}
let dates: Vec<String> = units
.iter()
.map(|u| u.start.format("%Y-%m-%d").to_string())
.collect();
assert_eq!(
dates,
["2026-03-07", "2026-03-08", "2026-03-09", "2026-03-10"]
);
let spring_forward = &units[1];
assert_eq!(
(spring_forward.end - spring_forward.start).num_hours(),
23,
"2026-03-08 loses an hour"
);
}
#[test]
fn calendar_day_windows_handle_fall_back() {
let ny = tz("America/New_York");
let from = parse_boundary("2026-10-31", ny).unwrap();
let to = parse_boundary("2026-11-03", ny).unwrap();
let units = plan_windows(from, to, Some(WindowStep::Days(1)), ny).unwrap();
for u in &units {
assert_eq!((u.start.hour(), u.start.minute()), (0, 0), "{}", u.id);
}
let long_day = units
.iter()
.find(|u| u.start.format("%Y-%m-%d").to_string() == "2026-11-01")
.expect("the fall-back day is planned");
assert_eq!((long_day.end - long_day.start).num_hours(), 25);
}
#[test]
fn calendar_and_absolute_windows_differ_across_dst() {
let ny = tz("America/New_York");
let from = parse_boundary("2026-03-07", ny).unwrap();
let to = parse_boundary("2026-03-10", ny).unwrap();
let cal = plan_windows(from, to, Some(parse_window("1d").unwrap()), ny).unwrap();
let abs = plan_windows(from, to, Some(parse_window("24h").unwrap()), ny).unwrap();
assert_eq!(cal[2].start.hour(), 0, "calendar stays on midnight");
assert_eq!(abs[2].start.hour(), 1, "absolute drifts by the DST delta");
assert_ne!(cal[2].start, abs[2].start);
}
#[test]
fn window_descriptor_is_stable_for_absolute_and_distinct_for_calendar() {
assert_eq!(
WindowStep::Absolute(Duration::hours(6)).to_string(),
"21600"
);
assert_eq!(WindowStep::Absolute(Duration::days(1)).to_string(), "86400");
assert_eq!(WindowStep::Days(1).to_string(), "1d");
assert_eq!(WindowStep::Weeks(2).to_string(), "2w");
}
#[test]
fn dst_transition_produces_no_gap_or_overlap() {
let ny = tz("America/New_York");
let from = parse_boundary("2026-03-07", ny).unwrap();
let to = parse_boundary("2026-03-10T00:00:00-04:00", ny).unwrap();
let units =
plan_windows(from, to, Some(WindowStep::Absolute(Duration::days(1))), ny).unwrap();
for w in units.windows(2) {
assert_eq!(w[0].end, w[1].start, "no gap/overlap across DST");
}
assert!(units[0].start.to_rfc3339().ends_with("-05:00"));
assert!(units.last().unwrap().end.to_rfc3339().ends_with("-04:00"));
}
#[test]
fn rejects_inverted_range_and_unit_explosion() {
let utc = tz("UTC");
let from = parse_boundary("2026-06-02", utc).unwrap();
let to = parse_boundary("2026-06-01", utc).unwrap();
assert!(plan_windows(from, to, None, utc).is_err());
let from = parse_boundary("2020-01-01", utc).unwrap();
let to = parse_boundary("2026-01-01", utc).unwrap();
let err = plan_windows(
from,
to,
Some(WindowStep::Absolute(Duration::minutes(1))),
utc,
)
.unwrap_err();
assert!(err.to_string().contains("larger window"), "{err}");
}
fn covered(chunks: &[IntChunk], bounds: Bounds) -> Vec<i64> {
let mut seen = Vec::new();
for c in chunks {
let last = match bounds {
Bounds::Inclusive => c.end,
Bounds::HalfOpen => c.end - 1,
};
for v in c.start..=last {
seen.push(v);
}
}
seen
}
#[test]
fn inclusive_chunks_tile_the_range_exactly_once() {
let chunks = plan_int_chunks(0, 24, 10, Bounds::Inclusive).unwrap();
assert_eq!(chunks.len(), 3);
assert_eq!((chunks[0].start, chunks[0].end), (0, 9));
assert_eq!((chunks[1].start, chunks[1].end), (10, 19));
assert_eq!((chunks[2].start, chunks[2].end), (20, 24), "last truncated");
assert_eq!(
covered(&chunks, Bounds::Inclusive),
(0..=24).collect::<Vec<_>>()
);
}
#[test]
fn half_open_chunks_tile_the_range_exactly_once() {
let chunks = plan_int_chunks(0, 25, 10, Bounds::HalfOpen).unwrap();
assert_eq!(chunks.len(), 3);
assert_eq!((chunks[0].start, chunks[0].end), (0, 10));
assert_eq!((chunks[1].start, chunks[1].end), (10, 20));
assert_eq!((chunks[2].start, chunks[2].end), (20, 25));
assert_eq!(
covered(&chunks, Bounds::HalfOpen),
(0..25).collect::<Vec<_>>()
);
}
#[test]
fn the_two_bounds_differ_by_exactly_one_at_every_boundary() {
let inc = plan_int_chunks(0, 19, 10, Bounds::Inclusive).unwrap();
let half = plan_int_chunks(0, 20, 10, Bounds::HalfOpen).unwrap();
assert_eq!(inc[0].end, 9);
assert_eq!(half[0].end, 10);
assert_eq!(inc[0].end + 1, half[0].end);
}
#[test]
fn tiles_exactly_once_across_many_sizes_and_ranges() {
for from in [-7i64, 0, 5, 1000] {
for span in [1i64, 2, 7, 10, 33, 100] {
for size in [1u64, 2, 3, 10, 64] {
let to = from + span - 1;
let chunks = plan_int_chunks(from, to, size, Bounds::Inclusive).unwrap();
assert_eq!(
covered(&chunks, Bounds::Inclusive),
(from..=to).collect::<Vec<_>>(),
"inclusive from={from} span={span} size={size}"
);
let chunks =
plan_int_chunks(from, from + span, size, Bounds::HalfOpen).unwrap();
assert_eq!(
covered(&chunks, Bounds::HalfOpen),
(from..from + span).collect::<Vec<_>>(),
"half-open from={from} span={span} size={size}"
);
}
}
}
}
#[test]
fn a_single_value_range_is_one_chunk_inclusive_and_empty_half_open() {
let inc = plan_int_chunks(5, 5, 10, Bounds::Inclusive).unwrap();
assert_eq!(inc.len(), 1);
assert_eq!((inc[0].start, inc[0].end), (5, 5));
assert!(plan_int_chunks(5, 5, 10, Bounds::HalfOpen).is_err());
}
#[test]
fn only_the_final_chunk_is_marked_last() {
let chunks = plan_int_chunks(0, 29, 10, Bounds::Inclusive).unwrap();
assert_eq!(
chunks.iter().filter(|c| c.is_last).count(),
1,
"exactly one chunk carries the open-ended tail flag"
);
assert!(chunks.last().unwrap().is_last);
}
#[test]
fn ids_are_zero_padded_so_they_sort_in_plan_order() {
let chunks = plan_int_chunks(0, 99, 1, Bounds::Inclusive).unwrap();
let mut ids: Vec<&str> = chunks.iter().map(|c| c.id.as_str()).collect();
let planned = ids.clone();
ids.sort_unstable();
assert_eq!(ids, planned, "lexicographic order must match plan order");
}
#[test]
fn rejects_inverted_and_empty_ranges() {
assert!(plan_int_chunks(10, 5, 10, Bounds::Inclusive).is_err());
assert!(plan_int_chunks(10, 10, 10, Bounds::HalfOpen).is_err());
}
#[test]
fn rejects_zero_chunk_size() {
let err = plan_int_chunks(0, 10, 0, Bounds::Inclusive).unwrap_err();
assert!(err.to_string().contains("greater than 0"), "{err}");
}
#[test]
fn rejects_a_chunk_explosion() {
let err = plan_int_chunks(0, 10_000_000, 1, Bounds::Inclusive).unwrap_err();
let msg = err.to_string();
assert!(msg.contains("larger chunk_size"), "{msg}");
assert!(msg.contains(&MAX_UNITS.to_string()), "names the cap: {msg}");
}
#[test]
fn does_not_overflow_near_i64_bounds() {
let chunks = plan_int_chunks(i64::MAX - 5, i64::MAX, 2, Bounds::Inclusive).unwrap();
assert_eq!(covered(&chunks, Bounds::Inclusive).len(), 6);
let chunks = plan_int_chunks(i64::MIN, i64::MIN + 5, 2, Bounds::Inclusive).unwrap();
assert_eq!(covered(&chunks, Bounds::Inclusive).len(), 6);
}
#[test]
fn offset_chunks_cover_the_total_without_overrunning_it() {
let chunks = plan_offset_chunks(25, 10).unwrap();
assert_eq!(chunks.len(), 3);
assert_eq!((chunks[0].offset, chunks[0].limit), (0, 10));
assert_eq!((chunks[1].offset, chunks[1].limit), (10, 10));
assert_eq!(
(chunks[2].offset, chunks[2].limit),
(20, 5),
"final limit is trimmed to the remainder"
);
assert_eq!(chunks.iter().map(|c| c.limit).sum::<u64>(), 25);
}
#[test]
fn an_exact_multiple_produces_full_chunks() {
let chunks = plan_offset_chunks(30, 10).unwrap();
assert_eq!(chunks.len(), 3);
assert!(chunks.iter().all(|c| c.limit == 10));
}
#[test]
fn a_zero_total_plans_nothing() {
assert!(plan_offset_chunks(0, 10).unwrap().is_empty());
}
#[test]
fn offset_rejects_zero_chunk_size_and_explosions() {
assert!(plan_offset_chunks(10, 0).is_err());
assert!(plan_offset_chunks(10_000_000, 1).is_err());
}
}