Skip to main content

garmin_cli/storage/
partitions.rs

1//! Partition key calculation for time-partitioned Parquet storage
2
3use chrono::{Datelike, NaiveDate};
4
5/// Entity types with their partition strategies
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum EntityType {
8    /// Weekly partitions (YYYY-Www)
9    Activities,
10    /// Daily partitions (YYYY-MM-DD)
11    TrackPoints,
12    /// Monthly partitions (YYYY-MM)
13    DailyHealth,
14    /// Monthly partitions (YYYY-MM)
15    PerformanceMetrics,
16    /// Monthly partitions (YYYY-MM)
17    Weight,
18    /// Single file (profiles.parquet)
19    Profiles,
20}
21
22impl EntityType {
23    /// Get the directory name for this entity
24    pub fn dir_name(&self) -> &'static str {
25        match self {
26            EntityType::Activities => "activities",
27            EntityType::TrackPoints => "track_points",
28            EntityType::DailyHealth => "daily_health",
29            EntityType::PerformanceMetrics => "performance_metrics",
30            EntityType::Weight => "weight",
31            EntityType::Profiles => "", // Single file, no directory
32        }
33    }
34
35    /// Calculate partition key for a given date
36    pub fn partition_key(&self, date: NaiveDate) -> String {
37        match self {
38            EntityType::Activities => {
39                // Weekly: YYYY-Www (ISO week)
40                format!("{}-W{:02}", date.iso_week().year(), date.iso_week().week())
41            }
42            EntityType::TrackPoints => {
43                // Daily: YYYY-MM-DD
44                date.format("%Y-%m-%d").to_string()
45            }
46            EntityType::DailyHealth | EntityType::PerformanceMetrics | EntityType::Weight => {
47                // Monthly: YYYY-MM
48                date.format("%Y-%m").to_string()
49            }
50            EntityType::Profiles => {
51                // Single file
52                "profiles".to_string()
53            }
54        }
55    }
56
57    /// Get the glob pattern for querying all partitions
58    pub fn glob_pattern(&self) -> String {
59        match self {
60            EntityType::Profiles => "profiles.parquet".to_string(),
61            _ => format!("{}/*.parquet", self.dir_name()),
62        }
63    }
64
65    /// Get the glob pattern for querying partitions in a date range
66    pub fn date_range_pattern(&self, from: NaiveDate, to: NaiveDate) -> Vec<String> {
67        let mut patterns = Vec::new();
68        let mut current = from;
69
70        while current <= to {
71            let key = self.partition_key(current);
72            let pattern = match self {
73                EntityType::Profiles => "profiles.parquet".to_string(),
74                _ => format!("{}/{}.parquet", self.dir_name(), key),
75            };
76
77            if !patterns.contains(&pattern) {
78                patterns.push(pattern);
79            }
80
81            // Advance by appropriate interval
82            current = match self {
83                EntityType::TrackPoints => current.succ_opt().unwrap_or(current),
84                EntityType::Activities => {
85                    // Advance to next week
86                    current + chrono::Duration::days(7)
87                }
88                _ => {
89                    // Advance to next month
90                    if current.month() == 12 {
91                        NaiveDate::from_ymd_opt(current.year() + 1, 1, 1).unwrap()
92                    } else {
93                        NaiveDate::from_ymd_opt(current.year(), current.month() + 1, 1).unwrap()
94                    }
95                }
96            };
97        }
98
99        patterns
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn test_weekly_partition_key() {
109        let date = NaiveDate::from_ymd_opt(2024, 12, 15).unwrap(); // Sunday of week 50
110        assert_eq!(EntityType::Activities.partition_key(date), "2024-W50");
111    }
112
113    #[test]
114    fn test_daily_partition_key() {
115        let date = NaiveDate::from_ymd_opt(2024, 12, 15).unwrap();
116        assert_eq!(EntityType::TrackPoints.partition_key(date), "2024-12-15");
117    }
118
119    #[test]
120    fn test_monthly_partition_key() {
121        let date = NaiveDate::from_ymd_opt(2024, 12, 15).unwrap();
122        assert_eq!(EntityType::DailyHealth.partition_key(date), "2024-12");
123    }
124
125    #[test]
126    fn test_glob_patterns() {
127        assert_eq!(
128            EntityType::Activities.glob_pattern(),
129            "activities/*.parquet"
130        );
131        assert_eq!(EntityType::Profiles.glob_pattern(), "profiles.parquet");
132    }
133}