datasynth_core/distributions/
timezone.rs1use chrono::{DateTime, NaiveDateTime, Offset, TimeZone, Utc};
9use chrono_tz::Tz;
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct TimezoneConfig {
16 pub default_timezone: String,
18 pub entity_timezones: HashMap<String, String>,
20 pub consolidation_timezone: String,
22}
23
24impl Default for TimezoneConfig {
25 fn default() -> Self {
26 Self {
27 default_timezone: "America/New_York".to_string(),
28 entity_timezones: HashMap::new(),
29 consolidation_timezone: "UTC".to_string(),
30 }
31 }
32}
33
34impl TimezoneConfig {
35 pub fn new(default_tz: &str) -> Self {
37 Self {
38 default_timezone: default_tz.to_string(),
39 entity_timezones: HashMap::new(),
40 consolidation_timezone: "UTC".to_string(),
41 }
42 }
43
44 pub fn with_consolidation(mut self, tz: &str) -> Self {
46 self.consolidation_timezone = tz.to_string();
47 self
48 }
49
50 pub fn add_mapping(mut self, entity_pattern: &str, timezone: &str) -> Self {
57 self.entity_timezones
58 .insert(entity_pattern.to_string(), timezone.to_string());
59 self
60 }
61}
62
63#[derive(Debug, Clone)]
65pub struct TimezoneHandler {
66 config: TimezoneConfig,
67 default_tz: Tz,
69 consolidation_tz: Tz,
71 entity_tz_cache: HashMap<String, Tz>,
73}
74
75impl TimezoneHandler {
76 pub fn new(config: TimezoneConfig) -> Result<Self, TimezoneError> {
78 let default_tz: Tz = config
79 .default_timezone
80 .parse()
81 .map_err(|_| TimezoneError::InvalidTimezone(config.default_timezone.clone()))?;
82
83 let consolidation_tz: Tz = config
84 .consolidation_timezone
85 .parse()
86 .map_err(|_| TimezoneError::InvalidTimezone(config.consolidation_timezone.clone()))?;
87
88 let mut entity_tz_cache = HashMap::new();
90 for (pattern, tz_name) in &config.entity_timezones {
91 let tz: Tz = tz_name
92 .parse()
93 .map_err(|_| TimezoneError::InvalidTimezone(tz_name.clone()))?;
94 entity_tz_cache.insert(pattern.clone(), tz);
95 }
96
97 Ok(Self {
98 config,
99 default_tz,
100 consolidation_tz,
101 entity_tz_cache,
102 })
103 }
104
105 pub fn us_eastern() -> Self {
107 Self::new(TimezoneConfig::default()).expect("Default timezone config should be valid")
108 }
109
110 pub fn get_entity_timezone(&self, entity_code: &str) -> Tz {
118 if let Some(tz) = self.entity_tz_cache.get(entity_code) {
120 return *tz;
121 }
122
123 for (pattern, tz) in &self.entity_tz_cache {
125 if let Some(prefix) = pattern.strip_suffix('*') {
126 if entity_code.starts_with(prefix) {
128 return *tz;
129 }
130 } else if let Some(suffix) = pattern.strip_prefix('*') {
131 if entity_code.ends_with(suffix) {
133 return *tz;
134 }
135 }
136 }
137
138 self.default_tz
139 }
140
141 pub fn to_utc(&self, local: NaiveDateTime, entity_code: &str) -> DateTime<Utc> {
143 let tz = self.get_entity_timezone(entity_code);
144 tz.from_local_datetime(&local)
145 .single()
146 .unwrap_or_else(|| {
147 tz.from_local_datetime(&local)
148 .earliest()
149 .expect("valid time components")
150 })
151 .with_timezone(&Utc)
152 }
153
154 pub fn to_local(&self, utc: DateTime<Utc>, entity_code: &str) -> NaiveDateTime {
156 let tz = self.get_entity_timezone(entity_code);
157 utc.with_timezone(&tz).naive_local()
158 }
159
160 pub fn to_consolidation(&self, local: NaiveDateTime, entity_code: &str) -> DateTime<Tz> {
162 let utc = self.to_utc(local, entity_code);
163 utc.with_timezone(&self.consolidation_tz)
164 }
165
166 pub fn consolidation_timezone(&self) -> Tz {
168 self.consolidation_tz
169 }
170
171 pub fn default_timezone(&self) -> Tz {
173 self.default_tz
174 }
175
176 pub fn get_timezone_name(&self, entity_code: &str) -> String {
178 self.get_entity_timezone(entity_code).name().to_string()
179 }
180
181 pub fn get_utc_offset_hours(&self, entity_code: &str, at: NaiveDateTime) -> f64 {
183 let tz = self.get_entity_timezone(entity_code);
184 let offset = tz.offset_from_local_datetime(&at).single();
185 match offset {
186 Some(o) => o.fix().local_minus_utc() as f64 / 3600.0,
187 None => 0.0,
188 }
189 }
190
191 pub fn config(&self) -> &TimezoneConfig {
193 &self.config
194 }
195}
196
197#[derive(Debug, Clone)]
199pub enum TimezoneError {
200 InvalidTimezone(String),
202 AmbiguousTime(NaiveDateTime),
204}
205
206impl std::fmt::Display for TimezoneError {
207 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
208 match self {
209 TimezoneError::InvalidTimezone(tz) => {
210 write!(f, "Invalid timezone: '{}'. Use IANA timezone names.", tz)
211 }
212 TimezoneError::AmbiguousTime(dt) => {
213 write!(f, "Ambiguous local time: {}", dt)
214 }
215 }
216 }
217}
218
219impl std::error::Error for TimezoneError {}
220
221pub struct TimezonePresets;
223
224impl TimezonePresets {
225 pub fn us_centric() -> TimezoneConfig {
227 TimezoneConfig::new("America/New_York")
228 .with_consolidation("America/New_York")
229 .add_mapping("*_WEST", "America/Los_Angeles")
230 .add_mapping("*_CENTRAL", "America/Chicago")
231 .add_mapping("*_MOUNTAIN", "America/Denver")
232 }
233
234 pub fn eu_centric() -> TimezoneConfig {
236 TimezoneConfig::new("Europe/London")
237 .with_consolidation("Europe/London")
238 .add_mapping("DE_*", "Europe/Berlin")
239 .add_mapping("FR_*", "Europe/Paris")
240 .add_mapping("CH_*", "Europe/Zurich")
241 }
242
243 pub fn apac_centric() -> TimezoneConfig {
245 TimezoneConfig::new("Asia/Singapore")
246 .with_consolidation("Asia/Singapore")
247 .add_mapping("JP_*", "Asia/Tokyo")
248 .add_mapping("CN_*", "Asia/Shanghai")
249 .add_mapping("IN_*", "Asia/Kolkata")
250 .add_mapping("AU_*", "Australia/Sydney")
251 }
252
253 pub fn global_utc() -> TimezoneConfig {
255 TimezoneConfig::new("America/New_York")
256 .with_consolidation("UTC")
257 .add_mapping("US_*", "America/New_York")
258 .add_mapping("EU_*", "Europe/London")
259 .add_mapping("APAC_*", "Asia/Singapore")
260 }
261}
262
263#[cfg(test)]
264#[allow(clippy::unwrap_used)]
265mod tests {
266 use super::*;
267 use chrono::Timelike;
268
269 #[test]
270 fn test_default_timezone() {
271 let handler = TimezoneHandler::us_eastern();
272 let tz = handler.get_entity_timezone("UNKNOWN_COMPANY");
273 assert_eq!(tz.name(), "America/New_York");
274 }
275
276 #[test]
277 fn test_exact_match() {
278 let config = TimezoneConfig::new("America/New_York").add_mapping("1000", "Europe/London");
279 let handler = TimezoneHandler::new(config).unwrap();
280
281 assert_eq!(handler.get_entity_timezone("1000").name(), "Europe/London");
282 assert_eq!(
283 handler.get_entity_timezone("2000").name(),
284 "America/New_York"
285 );
286 }
287
288 #[test]
289 fn test_prefix_pattern() {
290 let config = TimezoneConfig::new("America/New_York").add_mapping("EU_*", "Europe/Berlin");
291 let handler = TimezoneHandler::new(config).unwrap();
292
293 assert_eq!(
294 handler.get_entity_timezone("EU_1000").name(),
295 "Europe/Berlin"
296 );
297 assert_eq!(
298 handler.get_entity_timezone("EU_SUBSIDIARY").name(),
299 "Europe/Berlin"
300 );
301 assert_eq!(
302 handler.get_entity_timezone("US_1000").name(),
303 "America/New_York"
304 );
305 }
306
307 #[test]
308 fn test_suffix_pattern() {
309 let config = TimezoneConfig::new("America/New_York").add_mapping("*_APAC", "Asia/Tokyo");
310 let handler = TimezoneHandler::new(config).unwrap();
311
312 assert_eq!(
313 handler.get_entity_timezone("1000_APAC").name(),
314 "Asia/Tokyo"
315 );
316 assert_eq!(
317 handler.get_entity_timezone("CORP_APAC").name(),
318 "Asia/Tokyo"
319 );
320 assert_eq!(
321 handler.get_entity_timezone("1000_US").name(),
322 "America/New_York"
323 );
324 }
325
326 #[test]
327 fn test_to_utc() {
328 let handler = TimezoneHandler::new(TimezoneConfig::new("America/New_York")).unwrap();
329
330 let local =
332 NaiveDateTime::parse_from_str("2024-06-15 10:00:00", "%Y-%m-%d %H:%M:%S").unwrap();
333 let utc = handler.to_utc(local, "US_1000");
334
335 assert_eq!(utc.hour(), 14);
337 }
338
339 #[test]
340 fn test_to_local() {
341 let config = TimezoneConfig::new("America/New_York").add_mapping("EU_*", "Europe/London");
342 let handler = TimezoneHandler::new(config).unwrap();
343
344 let utc = DateTime::parse_from_rfc3339("2024-06-15T12:00:00Z")
346 .unwrap()
347 .with_timezone(&Utc);
348
349 let london_local = handler.to_local(utc, "EU_1000");
351 assert_eq!(london_local.hour(), 13);
352
353 let ny_local = handler.to_local(utc, "US_1000");
355 assert_eq!(ny_local.hour(), 8);
356 }
357
358 #[test]
359 fn test_presets() {
360 let _ = TimezoneHandler::new(TimezonePresets::us_centric()).unwrap();
362 let _ = TimezoneHandler::new(TimezonePresets::eu_centric()).unwrap();
363 let _ = TimezoneHandler::new(TimezonePresets::apac_centric()).unwrap();
364 let _ = TimezoneHandler::new(TimezonePresets::global_utc()).unwrap();
365 }
366
367 #[test]
368 fn test_invalid_timezone() {
369 let config = TimezoneConfig::new("Invalid/Timezone");
370 let result = TimezoneHandler::new(config);
371 assert!(result.is_err());
372 }
373}