servers/postgres/types/
interval.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Display;
16
17use bytes::{Buf, BufMut};
18use common_time::interval::IntervalFormat;
19use common_time::timestamp::TimeUnit;
20use common_time::{Duration, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
21use pgwire::types::format::FormatOptions;
22use pgwire::types::{FromSqlText, ToSqlText};
23use postgres_types::{FromSql, IsNull, ToSql, Type, to_sql_checked};
24
25use crate::error;
26
27/// On average one month has 30.44 day, which is a common approximation.
28const SECONDS_PER_MONTH: i64 = 24 * 6 * 6 * 3044;
29const SECONDS_PER_DAY: i64 = 24 * 60 * 60;
30const MILLISECONDS_PER_MONTH: i64 = SECONDS_PER_MONTH * 1000;
31const MILLISECONDS_PER_DAY: i64 = SECONDS_PER_DAY * 1000;
32
33#[derive(Debug, Clone, Copy, Default)]
34pub struct PgInterval {
35    pub(crate) months: i32,
36    pub(crate) days: i32,
37    pub(crate) microseconds: i64,
38}
39
40impl From<IntervalYearMonth> for PgInterval {
41    fn from(interval: IntervalYearMonth) -> Self {
42        Self {
43            months: interval.months,
44            days: 0,
45            microseconds: 0,
46        }
47    }
48}
49
50impl From<IntervalDayTime> for PgInterval {
51    fn from(interval: IntervalDayTime) -> Self {
52        Self {
53            months: 0,
54            days: interval.days,
55            microseconds: interval.milliseconds as i64 * 1000,
56        }
57    }
58}
59
60impl From<IntervalMonthDayNano> for PgInterval {
61    fn from(interval: IntervalMonthDayNano) -> Self {
62        Self {
63            months: interval.months,
64            days: interval.days,
65            microseconds: interval.nanoseconds / 1000,
66        }
67    }
68}
69
70impl TryFrom<Duration> for PgInterval {
71    type Error = error::Error;
72
73    fn try_from(duration: Duration) -> error::Result<Self> {
74        let value = duration.value();
75        let unit = duration.unit();
76
77        // Convert the duration to microseconds
78        match unit {
79            TimeUnit::Second => {
80                let months = i32::try_from(value / SECONDS_PER_MONTH)
81                    .map_err(|_| error::DurationOverflowSnafu { val: duration }.build())?;
82                let days =
83                    i32::try_from((value - (months as i64) * SECONDS_PER_MONTH) / SECONDS_PER_DAY)
84                        .map_err(|_| error::DurationOverflowSnafu { val: duration }.build())?;
85                let microseconds =
86                    (value - (months as i64) * SECONDS_PER_MONTH - (days as i64) * SECONDS_PER_DAY)
87                        .checked_mul(1_000_000)
88                        .ok_or(error::DurationOverflowSnafu { val: duration }.build())?;
89
90                Ok(Self {
91                    months,
92                    days,
93                    microseconds,
94                })
95            }
96            TimeUnit::Millisecond => {
97                let months = i32::try_from(value / MILLISECONDS_PER_MONTH)
98                    .map_err(|_| error::DurationOverflowSnafu { val: duration }.build())?;
99                let days = i32::try_from(
100                    (value - (months as i64) * MILLISECONDS_PER_MONTH) / MILLISECONDS_PER_DAY,
101                )
102                .map_err(|_| error::DurationOverflowSnafu { val: duration }.build())?;
103                let microseconds = ((value - (months as i64) * MILLISECONDS_PER_MONTH)
104                    - (days as i64) * MILLISECONDS_PER_DAY)
105                    * 1_000;
106                Ok(Self {
107                    months,
108                    days,
109                    microseconds,
110                })
111            }
112            TimeUnit::Microsecond => Ok(Self {
113                months: 0,
114                days: 0,
115                microseconds: value,
116            }),
117            TimeUnit::Nanosecond => Ok(Self {
118                months: 0,
119                days: 0,
120                microseconds: value / 1000,
121            }),
122        }
123    }
124}
125
126impl From<PgInterval> for IntervalMonthDayNano {
127    fn from(interval: PgInterval) -> Self {
128        IntervalMonthDayNano::new(
129            interval.months,
130            interval.days,
131            // Maybe overflow, but most scenarios ok.
132            interval.microseconds.checked_mul(1000).unwrap_or_else(|| {
133                if interval.microseconds.is_negative() {
134                    i64::MIN
135                } else {
136                    i64::MAX
137                }
138            }),
139        )
140    }
141}
142
143impl Display for PgInterval {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        write!(
146            f,
147            "{}",
148            IntervalFormat::from(IntervalMonthDayNano::from(*self)).to_postgres_string()
149        )
150    }
151}
152
153impl ToSql for PgInterval {
154    to_sql_checked!();
155
156    fn to_sql(
157        &self,
158        _: &Type,
159        out: &mut bytes::BytesMut,
160    ) -> std::result::Result<postgres_types::IsNull, Box<dyn snafu::Error + Sync + Send>>
161    where
162        Self: Sized,
163    {
164        // https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/timestamp.c#L989-L991
165        out.put_i64(self.microseconds);
166        out.put_i32(self.days);
167        out.put_i32(self.months);
168        Ok(postgres_types::IsNull::No)
169    }
170
171    fn accepts(ty: &Type) -> bool
172    where
173        Self: Sized,
174    {
175        matches!(ty, &Type::INTERVAL)
176    }
177}
178
179impl<'a> FromSql<'a> for PgInterval {
180    fn from_sql(
181        _: &Type,
182        mut raw: &'a [u8],
183    ) -> std::result::Result<Self, Box<dyn snafu::Error + Sync + Send>> {
184        // https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/timestamp.c#L1007-L1010
185        let microseconds = raw.get_i64();
186        let days = raw.get_i32();
187        let months = raw.get_i32();
188        Ok(PgInterval {
189            months,
190            days,
191            microseconds,
192        })
193    }
194
195    fn accepts(ty: &Type) -> bool {
196        matches!(ty, &Type::INTERVAL)
197    }
198}
199
200impl ToSqlText for PgInterval {
201    fn to_sql_text(
202        &self,
203        ty: &Type,
204        out: &mut bytes::BytesMut,
205        _format_options: &FormatOptions,
206    ) -> std::result::Result<postgres_types::IsNull, Box<dyn snafu::Error + Sync + Send>>
207    where
208        Self: Sized,
209    {
210        let fmt = match ty {
211            &Type::INTERVAL => self.to_string(),
212            _ => return Err("unsupported type".into()),
213        };
214
215        out.put_slice(fmt.as_bytes());
216        Ok(IsNull::No)
217    }
218}
219
220impl<'a> FromSqlText<'a> for PgInterval {
221    fn from_sql_text(
222        _ty: &Type,
223        input: &[u8],
224        _format_options: &FormatOptions,
225    ) -> std::result::Result<Self, Box<dyn snafu::Error + Sync + Send>>
226    where
227        Self: Sized,
228    {
229        // only support parsing interval from postgres format
230        if let Ok(interval) = pg_interval::Interval::from_postgres(str::from_utf8(input)?) {
231            Ok(PgInterval {
232                months: interval.months,
233                days: interval.days,
234                microseconds: interval.microseconds,
235            })
236        } else {
237            Err("invalid interval format".into())
238        }
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use common_time::Duration;
245    use common_time::timestamp::TimeUnit;
246
247    use super::*;
248
249    #[test]
250    fn test_duration_to_pg_interval() {
251        // Test with seconds
252        let duration = Duration::new(86400, TimeUnit::Second); // 1 day
253        let interval = PgInterval::try_from(duration).unwrap();
254        assert_eq!(interval.months, 0);
255        assert_eq!(interval.days, 1);
256        assert_eq!(interval.microseconds, 0);
257
258        // Test with milliseconds
259        let duration = Duration::new(86400000, TimeUnit::Millisecond); // 1 day
260        let interval = PgInterval::try_from(duration).unwrap();
261        assert_eq!(interval.months, 0);
262        assert_eq!(interval.days, 1);
263        assert_eq!(interval.microseconds, 0);
264
265        // Test with microseconds
266        let duration = Duration::new(86400000000, TimeUnit::Microsecond); // 1 day
267        let interval = PgInterval::try_from(duration).unwrap();
268        assert_eq!(interval.months, 0);
269        assert_eq!(interval.days, 0);
270        assert_eq!(interval.microseconds, 86400000000);
271
272        // Test with nanoseconds
273        let duration = Duration::new(86400000000000, TimeUnit::Nanosecond); // 1 day
274        let interval = PgInterval::try_from(duration).unwrap();
275        assert_eq!(interval.months, 0);
276        assert_eq!(interval.days, 0);
277        assert_eq!(interval.microseconds, 86400000000);
278
279        // Test with partial day
280        let duration = Duration::new(43200, TimeUnit::Second); // 12 hours
281        let interval = PgInterval::try_from(duration).unwrap();
282        assert_eq!(interval.months, 0);
283        assert_eq!(interval.days, 0);
284        assert_eq!(interval.microseconds, 43_200_000_000); // 12 hours in microseconds
285
286        // Test with negative duration
287        let duration = Duration::new(-86400, TimeUnit::Second); // -1 day
288        let interval = PgInterval::try_from(duration).unwrap();
289        assert_eq!(interval.months, 0);
290        assert_eq!(interval.days, -1);
291        assert_eq!(interval.microseconds, 0);
292
293        // Test with multiple days
294        let duration = Duration::new(259200, TimeUnit::Second); // 3 days
295        let interval = PgInterval::try_from(duration).unwrap();
296        assert_eq!(interval.months, 0);
297        assert_eq!(interval.days, 3);
298        assert_eq!(interval.microseconds, 0);
299
300        // Test with small duration (less than a day)
301        let duration = Duration::new(3600, TimeUnit::Second); // 1 hour
302        let interval = PgInterval::try_from(duration).unwrap();
303        assert_eq!(interval.months, 0);
304        assert_eq!(interval.days, 0);
305        assert_eq!(interval.microseconds, 3600000000); // 1 hour in microseconds
306
307        // Test with very small duration
308        let duration = Duration::new(1, TimeUnit::Microsecond); // 1 microsecond
309        let interval = PgInterval::try_from(duration).unwrap();
310        assert_eq!(interval.months, 0);
311        assert_eq!(interval.days, 0);
312        assert_eq!(interval.microseconds, 1);
313
314        let duration = Duration::new(i64::MAX, TimeUnit::Second);
315        assert!(PgInterval::try_from(duration).is_err());
316
317        let duration = Duration::new(i64::MAX, TimeUnit::Millisecond);
318        assert!(PgInterval::try_from(duration).is_err());
319    }
320}