1use 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
27const 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 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 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 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 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 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 let duration = Duration::new(86400, TimeUnit::Second); 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 let duration = Duration::new(86400000, TimeUnit::Millisecond); 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 let duration = Duration::new(86400000000, TimeUnit::Microsecond); 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 let duration = Duration::new(86400000000000, TimeUnit::Nanosecond); 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 let duration = Duration::new(43200, TimeUnit::Second); 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); let duration = Duration::new(-86400, TimeUnit::Second); 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 let duration = Duration::new(259200, TimeUnit::Second); 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 let duration = Duration::new(3600, TimeUnit::Second); 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); let duration = Duration::new(1, TimeUnit::Microsecond); 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}