mz_proto/
chrono.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Custom [`proptest::strategy::Strategy`] implementations and Protobuf types
11//! for the [`chrono`] fields used in the codebase.
12//!
13//! See the [`proptest`] docs[^1] for an example.
14//!
15//! [^1]: <https://altsysrq.github.io/proptest-book/proptest-derive/modifiers.html#strategy>
16
17use 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<ProtoNaiveDate> for NaiveDate {
30    fn into_proto(&self) -> ProtoNaiveDate {
31        ProtoNaiveDate {
32            year: self.year(),
33            ordinal: self.ordinal(),
34        }
35    }
36
37    fn from_proto(proto: ProtoNaiveDate) -> Result<Self, TryFromProtoError> {
38        NaiveDate::from_yo_opt(proto.year, proto.ordinal).ok_or_else(|| {
39            TryFromProtoError::DateConversionError(format!(
40                "NaiveDate::from_yo_opt({},{}) failed",
41                proto.year, proto.ordinal
42            ))
43        })
44    }
45}
46
47impl RustType<ProtoNaiveTime> for NaiveTime {
48    fn into_proto(&self) -> ProtoNaiveTime {
49        ProtoNaiveTime {
50            secs: self.num_seconds_from_midnight(),
51            frac: self.nanosecond(),
52        }
53    }
54
55    fn from_proto(proto: ProtoNaiveTime) -> Result<Self, TryFromProtoError> {
56        NaiveTime::from_num_seconds_from_midnight_opt(proto.secs, proto.frac).ok_or_else(|| {
57            TryFromProtoError::DateConversionError(format!(
58                "NaiveTime::from_num_seconds_from_midnight_opt({},{}) failed",
59                proto.secs, proto.frac
60            ))
61        })
62    }
63}
64
65impl RustType<ProtoNaiveDateTime> for NaiveDateTime {
66    fn into_proto(&self) -> ProtoNaiveDateTime {
67        ProtoNaiveDateTime {
68            year: self.year(),
69            ordinal: self.ordinal(),
70            secs: self.num_seconds_from_midnight(),
71            frac: self.nanosecond(),
72        }
73    }
74
75    fn from_proto(proto: ProtoNaiveDateTime) -> Result<Self, TryFromProtoError> {
76        let date = NaiveDate::from_yo_opt(proto.year, proto.ordinal).ok_or_else(|| {
77            TryFromProtoError::DateConversionError(format!(
78                "NaiveDate::from_yo_opt({},{}) failed",
79                proto.year, proto.ordinal
80            ))
81        })?;
82
83        let time = NaiveTime::from_num_seconds_from_midnight_opt(proto.secs, proto.frac)
84            .ok_or_else(|| {
85                TryFromProtoError::DateConversionError(format!(
86                    "NaiveTime::from_num_seconds_from_midnight_opt({},{}) failed",
87                    proto.secs, proto.frac
88                ))
89            })?;
90
91        Ok(NaiveDateTime::new(date, time))
92    }
93}
94
95impl RustType<ProtoNaiveDateTime> for DateTime<Utc> {
96    fn into_proto(&self) -> ProtoNaiveDateTime {
97        self.naive_utc().into_proto()
98    }
99
100    fn from_proto(proto: ProtoNaiveDateTime) -> Result<Self, TryFromProtoError> {
101        Ok(DateTime::from_naive_utc_and_offset(
102            NaiveDateTime::from_proto(proto)?,
103            Utc,
104        ))
105    }
106}
107
108impl RustType<ProtoFixedOffset> for FixedOffset {
109    fn into_proto(&self) -> ProtoFixedOffset {
110        ProtoFixedOffset {
111            local_minus_utc: self.local_minus_utc(),
112        }
113    }
114
115    fn from_proto(proto: ProtoFixedOffset) -> Result<Self, TryFromProtoError> {
116        FixedOffset::east_opt(proto.local_minus_utc).ok_or_else(|| {
117            TryFromProtoError::DateConversionError(format!(
118                "FixedOffset::east_opt({}) failed.",
119                proto.local_minus_utc
120            ))
121        })
122    }
123}
124
125/// Encode a Tz as string representation. This is not the most space efficient solution, but
126/// it is immune to changes in the chrono_tz (and is fully compatible with its public API).
127impl RustType<ProtoTz> for chrono_tz::Tz {
128    fn into_proto(&self) -> ProtoTz {
129        ProtoTz {
130            name: self.name().into(),
131        }
132    }
133
134    fn from_proto(proto: ProtoTz) -> Result<Self, TryFromProtoError> {
135        Tz::from_str(&proto.name).map_err(TryFromProtoError::DateConversionError)
136    }
137}
138
139pub fn any_naive_date() -> impl Strategy<Value = NaiveDate> {
140    (0..1000000).prop_map(|d| NaiveDate::from_num_days_from_ce_opt(d).unwrap())
141}
142
143pub fn any_naive_datetime() -> impl Strategy<Value = NaiveDateTime> {
144    (0..(NaiveDateTime::MAX.nanosecond() - NaiveDateTime::MIN.nanosecond()))
145        .prop_map(|x| NaiveDateTime::MIN + Duration::nanoseconds(i64::from(x)))
146}
147
148pub fn any_datetime() -> impl Strategy<Value = DateTime<Utc>> {
149    any_naive_datetime().prop_map(|x| DateTime::from_naive_utc_and_offset(x, Utc))
150}
151
152pub fn any_fixed_offset() -> impl Strategy<Value = FixedOffset> {
153    (-86_399..86_400).prop_map(|o| FixedOffset::east_opt(o).unwrap())
154}
155
156pub fn any_timezone() -> impl Strategy<Value = Tz> {
157    (0..TZ_VARIANTS.len()).prop_map(|idx| *TZ_VARIANTS.get(idx).unwrap())
158}
159
160#[cfg(test)]
161mod tests {
162    use crate::protobuf_roundtrip;
163    use mz_ore::assert_ok;
164    use proptest::prelude::*;
165
166    use super::*;
167
168    proptest! {
169        #![proptest_config(ProptestConfig::with_cases(4096))]
170
171        #[mz_ore::test]
172        #[cfg_attr(miri, ignore)] // too slow
173        fn naive_date_protobuf_roundtrip(expect in any_naive_date() ) {
174            let actual = protobuf_roundtrip::<_, ProtoNaiveDate>(&expect);
175            assert_ok!(actual);
176            assert_eq!(actual.unwrap(), expect);
177        }
178
179        #[mz_ore::test]
180        #[cfg_attr(miri, ignore)] // too slow
181        fn naive_date_time_protobuf_roundtrip(expect in any_naive_datetime() ) {
182            let actual = protobuf_roundtrip::<_, ProtoNaiveDateTime>(&expect);
183            assert_ok!(actual);
184            assert_eq!(actual.unwrap(), expect);
185        }
186
187        #[mz_ore::test]
188        #[cfg_attr(miri, ignore)] // too slow
189        fn date_time_protobuf_roundtrip(expect in any_datetime() ) {
190            let actual = protobuf_roundtrip::<_, ProtoNaiveDateTime>(&expect);
191            assert_ok!(actual);
192            assert_eq!(actual.unwrap(), expect);
193        }
194
195        #[mz_ore::test]
196        #[cfg_attr(miri, ignore)] // too slow
197        fn fixed_offset_protobuf_roundtrip(expect in any_fixed_offset() ) {
198            let actual = protobuf_roundtrip::<_, ProtoFixedOffset>(&expect);
199            assert_ok!(actual);
200            assert_eq!(actual.unwrap(), expect);
201        }
202    }
203}