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<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
107/// Encode a Tz as string representation. This is not the most space efficient solution, but
108/// it is immune to changes in the chrono_tz (and is fully compatible with its public API).
109impl 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)] // too slow
155        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)] // too slow
163        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)] // too slow
171        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}