1use std::str::FromStr;
18
19use chrono::{
20 DateTime, Datelike, Duration, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, Timelike, Utc,
21};
22use chrono_tz::{TZ_VARIANTS, Tz};
23use proptest::prelude::Strategy;
24
25use crate::{RustType, TryFromProtoError};
26
27include!(concat!(env!("OUT_DIR"), "/mz_proto.chrono.rs"));
28
29impl RustType<ProtoNaiveTime> for NaiveTime {
30 fn into_proto(&self) -> ProtoNaiveTime {
31 ProtoNaiveTime {
32 secs: self.num_seconds_from_midnight(),
33 frac: self.nanosecond(),
34 }
35 }
36
37 fn from_proto(proto: ProtoNaiveTime) -> Result<Self, TryFromProtoError> {
38 NaiveTime::from_num_seconds_from_midnight_opt(proto.secs, proto.frac).ok_or_else(|| {
39 TryFromProtoError::DateConversionError(format!(
40 "NaiveTime::from_num_seconds_from_midnight_opt({},{}) failed",
41 proto.secs, proto.frac
42 ))
43 })
44 }
45}
46
47impl RustType<ProtoNaiveDateTime> for NaiveDateTime {
48 fn into_proto(&self) -> ProtoNaiveDateTime {
49 ProtoNaiveDateTime {
50 year: self.year(),
51 ordinal: self.ordinal(),
52 secs: self.num_seconds_from_midnight(),
53 frac: self.nanosecond(),
54 }
55 }
56
57 fn from_proto(proto: ProtoNaiveDateTime) -> Result<Self, TryFromProtoError> {
58 let date = NaiveDate::from_yo_opt(proto.year, proto.ordinal).ok_or_else(|| {
59 TryFromProtoError::DateConversionError(format!(
60 "NaiveDate::from_yo_opt({},{}) failed",
61 proto.year, proto.ordinal
62 ))
63 })?;
64
65 let time = NaiveTime::from_num_seconds_from_midnight_opt(proto.secs, proto.frac)
66 .ok_or_else(|| {
67 TryFromProtoError::DateConversionError(format!(
68 "NaiveTime::from_num_seconds_from_midnight_opt({},{}) failed",
69 proto.secs, proto.frac
70 ))
71 })?;
72
73 Ok(NaiveDateTime::new(date, time))
74 }
75}
76
77impl RustType<ProtoNaiveDateTime> for DateTime<Utc> {
78 fn into_proto(&self) -> ProtoNaiveDateTime {
79 self.naive_utc().into_proto()
80 }
81
82 fn from_proto(proto: ProtoNaiveDateTime) -> Result<Self, TryFromProtoError> {
83 Ok(DateTime::from_naive_utc_and_offset(
84 NaiveDateTime::from_proto(proto)?,
85 Utc,
86 ))
87 }
88}
89
90impl RustType<ProtoFixedOffset> for FixedOffset {
91 fn into_proto(&self) -> ProtoFixedOffset {
92 ProtoFixedOffset {
93 local_minus_utc: self.local_minus_utc(),
94 }
95 }
96
97 fn from_proto(proto: ProtoFixedOffset) -> Result<Self, TryFromProtoError> {
98 FixedOffset::east_opt(proto.local_minus_utc).ok_or_else(|| {
99 TryFromProtoError::DateConversionError(format!(
100 "FixedOffset::east_opt({}) failed.",
101 proto.local_minus_utc
102 ))
103 })
104 }
105}
106
107impl RustType<ProtoTz> for chrono_tz::Tz {
110 fn into_proto(&self) -> ProtoTz {
111 ProtoTz {
112 name: self.name().into(),
113 }
114 }
115
116 fn from_proto(proto: ProtoTz) -> Result<Self, TryFromProtoError> {
117 Tz::from_str(&proto.name).map_err(TryFromProtoError::DateConversionError)
118 }
119}
120
121pub fn any_naive_date() -> impl Strategy<Value = NaiveDate> {
122 (0..1000000).prop_map(|d| NaiveDate::from_num_days_from_ce_opt(d).unwrap())
123}
124
125pub fn any_naive_datetime() -> impl Strategy<Value = NaiveDateTime> {
126 (0..(NaiveDateTime::MAX.nanosecond() - NaiveDateTime::MIN.nanosecond()))
127 .prop_map(|x| NaiveDateTime::MIN + Duration::nanoseconds(i64::from(x)))
128}
129
130pub fn any_datetime() -> impl Strategy<Value = DateTime<Utc>> {
131 any_naive_datetime().prop_map(|x| DateTime::from_naive_utc_and_offset(x, Utc))
132}
133
134pub fn any_fixed_offset() -> impl Strategy<Value = FixedOffset> {
135 (-86_399..86_400).prop_map(|o| FixedOffset::east_opt(o).unwrap())
136}
137
138pub fn any_timezone() -> impl Strategy<Value = Tz> {
139 (0..TZ_VARIANTS.len()).prop_map(|idx| *TZ_VARIANTS.get(idx).unwrap())
140}
141
142#[cfg(test)]
143mod tests {
144 use crate::protobuf_roundtrip;
145 use mz_ore::assert_ok;
146 use proptest::prelude::*;
147
148 use super::*;
149
150 proptest! {
151 #![proptest_config(ProptestConfig::with_cases(4096))]
152
153 #[mz_ore::test]
154 #[cfg_attr(miri, ignore)] fn naive_date_time_protobuf_roundtrip(expect in any_naive_datetime() ) {
156 let actual = protobuf_roundtrip::<_, ProtoNaiveDateTime>(&expect);
157 assert_ok!(actual);
158 assert_eq!(actual.unwrap(), expect);
159 }
160
161 #[mz_ore::test]
162 #[cfg_attr(miri, ignore)] fn date_time_protobuf_roundtrip(expect in any_datetime() ) {
164 let actual = protobuf_roundtrip::<_, ProtoNaiveDateTime>(&expect);
165 assert_ok!(actual);
166 assert_eq!(actual.unwrap(), expect);
167 }
168
169 #[mz_ore::test]
170 #[cfg_attr(miri, ignore)] fn fixed_offset_protobuf_roundtrip(expect in any_fixed_offset() ) {
172 let actual = protobuf_roundtrip::<_, ProtoFixedOffset>(&expect);
173 assert_ok!(actual);
174 assert_eq!(actual.unwrap(), expect);
175 }
176 }
177}