risingwave_common/types/
timestamptz.rs1use std::error::Error;
16use std::io::{Cursor, Write};
17use std::str::FromStr;
18
19use anyhow::Context;
20use byteorder::{BigEndian, ReadBytesExt};
21use bytes::BytesMut;
22use chrono::{DateTime, Datelike, TimeZone, Utc};
23use chrono_tz::Tz;
24use postgres_types::{FromSql, IsNull, ToSql, Type, accepts, to_sql_checked};
25use risingwave_common_estimate_size::ZeroHeapSize;
26use serde::{Deserialize, Serialize};
27
28use super::DataType;
29use super::to_text::ToText;
30use crate::array::ArrayResult;
31
32#[derive(
34 Default, Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
35)]
36#[repr(transparent)]
37pub struct Timestamptz(i64);
38
39impl ZeroHeapSize for Timestamptz {}
40
41impl ToSql for Timestamptz {
42 accepts!(TIMESTAMPTZ);
43
44 to_sql_checked!();
45
46 fn to_sql(&self, _: &Type, out: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>>
47 where
48 Self: Sized,
49 {
50 let instant = self.to_datetime_utc();
51 instant.to_sql(&Type::ANY, out)
52 }
53}
54
55impl<'a> FromSql<'a> for Timestamptz {
56 fn from_sql(
57 ty: &Type,
58 raw: &'a [u8],
59 ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
60 let instant = DateTime::<Utc>::from_sql(ty, raw)?;
61 Ok(Self::from(instant))
62 }
63
64 fn accepts(ty: &Type) -> bool {
65 matches!(*ty, Type::TIMESTAMPTZ)
66 }
67}
68
69impl ToText for Timestamptz {
70 fn write<W: std::fmt::Write>(&self, f: &mut W) -> std::fmt::Result {
71 let instant = self.to_datetime_utc();
74 write!(f, "{}+00:00", instant.naive_local())
78 }
79
80 fn write_with_type<W: std::fmt::Write>(&self, ty: &DataType, f: &mut W) -> std::fmt::Result {
81 assert_eq!(ty, &DataType::Timestamptz);
82 self.write(f)
83 }
84}
85
86impl Timestamptz {
87 pub fn from_secs(timestamp_secs: i64) -> Option<Self> {
89 timestamp_secs
90 .checked_mul(1_000_000)
91 .and_then(Self::from_micros)
92 }
93
94 pub fn from_millis(timestamp_millis: i64) -> Option<Self> {
97 timestamp_millis
98 .checked_mul(1000)
99 .and_then(Self::from_micros)
100 }
101
102 pub fn from_micros(timestamp_micros: i64) -> Option<Self> {
105 const MIN_MICROS: i64 = DateTime::<Utc>::MIN_UTC.timestamp_micros();
107 const MAX_MICROS: i64 = DateTime::<Utc>::MAX_UTC.timestamp_micros();
108 (MIN_MICROS..=MAX_MICROS)
109 .contains(×tamp_micros)
110 .then_some(Self(timestamp_micros))
111 }
112
113 pub fn from_micros_uncheck(timestamp_micros: i64) -> Self {
118 Self(timestamp_micros)
119 }
120
121 pub fn from_nanos(timestamp_nanos: i64) -> Self {
126 Self(timestamp_nanos.div_euclid(1_000))
127 }
128
129 pub fn timestamp_micros(&self) -> i64 {
131 self.0
132 }
133
134 pub fn timestamp_millis(&self) -> i64 {
136 self.0.div_euclid(1_000)
137 }
138
139 pub fn timestamp_nanos(&self) -> Option<i64> {
141 self.0.checked_mul(1_000)
142 }
143
144 pub fn timestamp(&self) -> i64 {
147 self.0.div_euclid(1_000_000)
148 }
149
150 pub fn timestamp_subsec_nanos(&self) -> u32 {
152 self.0.rem_euclid(1_000_000) as u32 * 1000
153 }
154
155 pub fn to_datetime_utc(self) -> chrono::DateTime<Utc> {
156 self.into()
157 }
158
159 pub fn to_datetime_in_zone(self, tz: Tz) -> chrono::DateTime<Tz> {
160 self.to_datetime_utc().with_timezone(&tz)
161 }
162
163 pub fn lookup_time_zone(time_zone: &str) -> std::result::Result<Tz, String> {
164 Tz::from_str_insensitive(time_zone)
165 .map_err(|_| format!("'{time_zone}' is not a valid timezone"))
166 }
167
168 pub fn from_protobuf(cur: &mut Cursor<&[u8]>) -> ArrayResult<Timestamptz> {
169 let micros = cur
170 .read_i64::<BigEndian>()
171 .context("failed to read i64 from Timestamptz buffer")?;
172 Ok(Self(micros))
173 }
174
175 pub fn to_protobuf(self, output: &mut impl Write) -> ArrayResult<usize> {
176 output.write(&self.0.to_be_bytes()).map_err(Into::into)
177 }
178}
179
180impl<Tz: TimeZone> From<chrono::DateTime<Tz>> for Timestamptz {
181 fn from(dt: chrono::DateTime<Tz>) -> Self {
182 Self(dt.timestamp_micros())
183 }
184}
185
186impl From<Timestamptz> for chrono::DateTime<Utc> {
187 fn from(tz: Timestamptz) -> Self {
188 Utc.timestamp_opt(tz.timestamp(), tz.timestamp_subsec_nanos())
192 .unwrap()
193 }
194}
195
196impl FromStr for Timestamptz {
197 type Err = &'static str;
198
199 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
200 pub const ERROR_MSG: &str = concat!(
201 "Can't cast string to timestamp with time zone (expected format is YYYY-MM-DD HH:MM:SS[.D+{up to 6 digits}] followed by +hh:mm or literal Z)",
202 "\nFor example: '2021-04-01 00:00:00+00:00'"
203 );
204 let ret = match speedate::DateTime::parse_str_rfc3339(s) {
208 Ok(r) => r,
209 Err(_) => {
210 return s
213 .parse::<chrono::DateTime<Utc>>()
214 .or_else(|_| {
215 chrono::DateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f%#z")
216 .map(|t| t.with_timezone(&Utc))
217 })
218 .map(|t| Timestamptz(t.timestamp_micros()))
219 .map_err(|_| ERROR_MSG);
220 }
221 };
222 if ret.time.tz_offset.is_none() {
223 return Err(ERROR_MSG);
224 }
225 Ok(Timestamptz(
226 ret.timestamp_tz()
227 .checked_mul(1000000)
228 .and_then(|us| us.checked_add(ret.time.microsecond.into()))
229 .ok_or(ERROR_MSG)?,
230 ))
231 }
232}
233
234impl std::fmt::Display for Timestamptz {
235 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
236 self.write(f)
237 }
238}
239
240pub fn write_date_time_tz(
241 instant_local: DateTime<Tz>,
242 writer: &mut impl std::fmt::Write,
243) -> std::fmt::Result {
244 let date = instant_local.date_naive();
245 let (ce, year) = date.year_ce();
246 write!(
247 writer,
248 "{:04}-{:02}-{:02} {}",
249 year,
250 date.month(),
251 date.day(),
252 instant_local.format(if ce {
253 "%H:%M:%S%.f%:z"
254 } else {
255 "%H:%M:%S%.f%:z BC"
256 })
257 )
258}
259
260#[cfg(test)]
261mod test {
262 use super::*;
263
264 #[test]
265 fn from_micros_checks_representability() {
266 for micros in [
268 0,
269 i64::MAX,
270 i64::MIN,
271 8_210_266_876_799_999_999, 8_210_266_876_800_000_000, -8_334_601_228_800_000_000, -8_334_601_228_800_000_001, ] {
276 assert_eq!(
277 Timestamptz::from_micros(micros).is_some(),
278 DateTime::from_timestamp_micros(micros).is_some(),
279 "mismatch at {micros}"
280 );
281 }
282 }
283
284 #[test]
285 fn parse() {
286 assert!("1999-01-08 04:05:06".parse::<Timestamptz>().is_err());
287 assert_eq!(
288 "2022-08-03 10:34:02Z".parse::<Timestamptz>().unwrap(),
289 "2022-08-03 02:34:02-08:00".parse::<Timestamptz>().unwrap()
290 );
291
292 let expected = Ok(Timestamptz::from_micros(1689130892000000).unwrap());
293 assert_eq!("2023-07-12T03:01:32Z".parse(), expected);
295 assert_eq!("2023-07-12T03:01:32+00:00".parse(), expected);
296 assert_eq!("2023-07-12T11:01:32+08:00".parse(), expected);
297 assert_eq!("2023-07-12 03:01:32Z".parse(), expected);
299 assert_eq!("2023-07-12 03:01:32+00:00".parse(), expected);
300 assert_eq!("2023-07-12 11:01:32+08:00".parse(), expected);
301 assert_eq!("2023-07-12 03:01:32+00".parse(), expected);
303 assert_eq!("2023-07-12 11:01:32+08".parse(), expected);
304 }
305}