archiver_core/storage/
partition.rs

1use std::time::SystemTime;
2
3use chrono::{Datelike, Duration, NaiveDate, Timelike, Utc};
4use serde::{Deserialize, Serialize};
5
6/// Time-based partition granularity — matches Java archiver's PartitionGranularity.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
8pub enum PartitionGranularity {
9    #[serde(rename = "5min")]
10    FiveMin,
11    #[serde(rename = "15min")]
12    FifteenMin,
13    #[serde(rename = "30min")]
14    ThirtyMin,
15    #[serde(rename = "hour")]
16    Hour,
17    #[serde(rename = "day")]
18    Day,
19    #[serde(rename = "month")]
20    Month,
21    #[serde(rename = "year")]
22    Year,
23}
24
25impl PartitionGranularity {
26    pub fn approx_seconds(self) -> u64 {
27        match self {
28            Self::FiveMin => 5 * 60,
29            Self::FifteenMin => 15 * 60,
30            Self::ThirtyMin => 30 * 60,
31            Self::Hour => 3600,
32            Self::Day => 86400,
33            Self::Month => 31 * 86400,
34            Self::Year => 366 * 86400,
35        }
36    }
37
38    fn approx_minutes(self) -> u32 {
39        (self.approx_seconds() / 60) as u32
40    }
41}
42
43/// Generate the partition name string for a given timestamp.
44/// Matches Java archiver's TimeUtils.getPartitionName().
45///
46/// Examples:
47///   Year  → "2024"
48///   Month → "2024_03"
49///   Day   → "2024_03_15"
50///   Hour  → "2024_03_15_09"
51///   5Min  → "2024_03_15_09_30"
52pub fn partition_name(ts: SystemTime, granularity: PartitionGranularity) -> String {
53    let dt = chrono::DateTime::<Utc>::from(ts);
54    let y = dt.year();
55    let m = dt.month();
56    let d = dt.day();
57    let h = dt.hour();
58    let min = dt.minute();
59
60    match granularity {
61        PartitionGranularity::Year => format!("{y}"),
62        PartitionGranularity::Month => format!("{y}_{m:02}"),
63        PartitionGranularity::Day => format!("{y}_{m:02}_{d:02}"),
64        PartitionGranularity::Hour => format!("{y}_{m:02}_{d:02}_{h:02}"),
65        PartitionGranularity::FiveMin
66        | PartitionGranularity::FifteenMin
67        | PartitionGranularity::ThirtyMin => {
68            let approx_min = granularity.approx_minutes();
69            let start_min = (min / approx_min) * approx_min;
70            format!("{y}_{m:02}_{d:02}_{h:02}_{start_min:02}")
71        }
72    }
73}
74
75/// List all partition names that overlap with [start, end].
76pub fn partitions_in_range(
77    start: SystemTime,
78    end: SystemTime,
79    granularity: PartitionGranularity,
80) -> Vec<String> {
81    let mut names = Vec::new();
82    let mut current = start;
83    loop {
84        let name = partition_name(current, granularity);
85        if names.last().map(|n: &String| n.as_str()) != Some(&name) {
86            names.push(name);
87        }
88        if current >= end {
89            break;
90        }
91        current = next_partition_start(current, granularity);
92        if current > end {
93            // Include the last partition.
94            let name = partition_name(end, granularity);
95            if names.last().map(|n: &String| n.as_str()) != Some(&name) {
96                names.push(name);
97            }
98            break;
99        }
100    }
101    names
102}
103
104/// Get the most recent N partition names up to and including the one containing `ts`.
105pub fn recent_partitions(
106    ts: SystemTime,
107    granularity: PartitionGranularity,
108    count: usize,
109) -> Vec<String> {
110    let mut names = Vec::new();
111    let mut current = ts;
112    for _ in 0..count {
113        names.push(partition_name(current, granularity));
114        current = prev_partition_end(current, granularity);
115    }
116    names.reverse();
117    names
118}
119
120/// Compute the start of the next partition after the one containing `ts`.
121pub fn next_partition_start(ts: SystemTime, granularity: PartitionGranularity) -> SystemTime {
122    let dt = chrono::DateTime::<Utc>::from(ts);
123
124    let next = match granularity {
125        PartitionGranularity::Year => NaiveDate::from_ymd_opt(dt.year() + 1, 1, 1)
126            .expect("Jan 1 is always valid")
127            .and_hms_opt(0, 0, 0)
128            .expect("midnight is always valid")
129            .and_utc(),
130        PartitionGranularity::Month => {
131            let (y, m) = if dt.month() == 12 {
132                (dt.year() + 1, 1)
133            } else {
134                (dt.year(), dt.month() + 1)
135            };
136            NaiveDate::from_ymd_opt(y, m, 1)
137                .expect("1st of month is always valid")
138                .and_hms_opt(0, 0, 0)
139                .expect("midnight is always valid")
140                .and_utc()
141        }
142        PartitionGranularity::Day => (dt.date_naive() + Duration::days(1))
143            .and_hms_opt(0, 0, 0)
144            .expect("midnight is always valid")
145            .and_utc(),
146        PartitionGranularity::Hour => {
147            let current_hour = dt
148                .date_naive()
149                .and_hms_opt(dt.hour(), 0, 0)
150                .expect("hour from valid DateTime")
151                .and_utc();
152            current_hour + Duration::hours(1)
153        }
154        PartitionGranularity::FiveMin
155        | PartitionGranularity::FifteenMin
156        | PartitionGranularity::ThirtyMin => {
157            let approx_min = granularity.approx_minutes();
158            let start_min = (dt.minute() / approx_min) * approx_min;
159            let current_start = dt
160                .date_naive()
161                .and_hms_opt(dt.hour(), start_min, 0)
162                .expect("aligned minute from valid DateTime")
163                .and_utc();
164            current_start + Duration::minutes(approx_min as i64)
165        }
166    };
167
168    next.into()
169}
170
171/// Compute the last moment of the previous partition before `ts`.
172fn prev_partition_end(ts: SystemTime, granularity: PartitionGranularity) -> SystemTime {
173    let dt = chrono::DateTime::<Utc>::from(ts);
174
175    let prev_end = match granularity {
176        PartitionGranularity::Year => NaiveDate::from_ymd_opt(dt.year() - 1, 12, 31)
177            .expect("Dec 31 is always valid")
178            .and_hms_opt(23, 59, 59)
179            .expect("23:59:59 is always valid")
180            .and_utc(),
181        PartitionGranularity::Month => {
182            let first_of_month = NaiveDate::from_ymd_opt(dt.year(), dt.month(), 1)
183                .expect("1st of month is always valid");
184            let prev = first_of_month - Duration::days(1);
185            prev.and_hms_opt(23, 59, 59)
186                .expect("23:59:59 is always valid")
187                .and_utc()
188        }
189        PartitionGranularity::Day => {
190            let prev = dt.date_naive() - Duration::days(1);
191            prev.and_hms_opt(23, 59, 59)
192                .expect("23:59:59 is always valid")
193                .and_utc()
194        }
195        PartitionGranularity::Hour => {
196            dt.date_naive()
197                .and_hms_opt(dt.hour(), 0, 0)
198                .expect("hour from valid DateTime")
199                .and_utc()
200                - Duration::seconds(1)
201        }
202        PartitionGranularity::FiveMin
203        | PartitionGranularity::FifteenMin
204        | PartitionGranularity::ThirtyMin => {
205            let approx_min = granularity.approx_minutes();
206            let start_min = (dt.minute() / approx_min) * approx_min;
207            dt.date_naive()
208                .and_hms_opt(dt.hour(), start_min, 0)
209                .expect("aligned minute from valid DateTime")
210                .and_utc()
211                - Duration::seconds(1)
212        }
213    };
214
215    prev_end.into()
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221    use chrono::TimeZone;
222
223    #[test]
224    fn test_partition_name_year() {
225        let ts: SystemTime = Utc.with_ymd_and_hms(2024, 6, 15, 10, 30, 0).unwrap().into();
226        assert_eq!(partition_name(ts, PartitionGranularity::Year), "2024");
227    }
228
229    #[test]
230    fn test_partition_name_month() {
231        let ts: SystemTime = Utc.with_ymd_and_hms(2024, 3, 15, 10, 30, 0).unwrap().into();
232        assert_eq!(partition_name(ts, PartitionGranularity::Month), "2024_03");
233    }
234
235    #[test]
236    fn test_partition_name_day() {
237        let ts: SystemTime = Utc.with_ymd_and_hms(2024, 3, 5, 10, 30, 0).unwrap().into();
238        assert_eq!(partition_name(ts, PartitionGranularity::Day), "2024_03_05");
239    }
240
241    #[test]
242    fn test_partition_name_hour() {
243        let ts: SystemTime = Utc.with_ymd_and_hms(2024, 3, 5, 9, 30, 0).unwrap().into();
244        assert_eq!(
245            partition_name(ts, PartitionGranularity::Hour),
246            "2024_03_05_09"
247        );
248    }
249
250    #[test]
251    fn test_partition_name_15min() {
252        let ts: SystemTime = Utc.with_ymd_and_hms(2024, 3, 5, 9, 47, 0).unwrap().into();
253        assert_eq!(
254            partition_name(ts, PartitionGranularity::FifteenMin),
255            "2024_03_05_09_45"
256        );
257    }
258
259    #[test]
260    fn test_partitions_in_range() {
261        let start: SystemTime = Utc.with_ymd_and_hms(2024, 3, 5, 10, 0, 0).unwrap().into();
262        let end: SystemTime = Utc.with_ymd_and_hms(2024, 3, 5, 12, 30, 0).unwrap().into();
263        let names = partitions_in_range(start, end, PartitionGranularity::Hour);
264        assert_eq!(
265            names,
266            vec!["2024_03_05_10", "2024_03_05_11", "2024_03_05_12"]
267        );
268    }
269}
archiver_core/storage/partition.rs

archiver_core/storage/
partition.rs
