Skip to main content

risingwave_common/types/
timestamptz.rs

1// Copyright 2023 RisingWave Labs
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::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/// Timestamp with timezone.
33#[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        // Just a meaningful representation as placeholder. The real implementation depends
72        // on TimeZone from session. See #3552.
73        let instant = self.to_datetime_utc();
74        // PostgreSQL uses a space rather than `T` to separate the date and time.
75        // https://www.postgresql.org/docs/current/datatype-datetime.html#DATATYPE-DATETIME-OUTPUT
76        // same as `instant.format("%Y-%m-%d %H:%M:%S%.f%:z")` but faster
77        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    /// Creates a `Timestamptz` from seconds. Returns `None` if the given timestamp is out of range.
88    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    /// Creates a `Timestamptz` from milliseconds. Returns `None` if the given timestamp is out of
95    /// range.
96    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    /// Creates a `Timestamptz` from microseconds. Returns `None` if the value does not convert to
103    /// [`chrono::DateTime`], which formatting and time zone operations require.
104    pub fn from_micros(timestamp_micros: i64) -> Option<Self> {
105        // Exactly the values `DateTime::from_timestamp_micros` accepts, as a cheap range check.
106        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(&timestamp_micros)
110            .then_some(Self(timestamp_micros))
111    }
112
113    /// Creates a `Timestamptz` from microseconds without checking representability.
114    ///
115    /// For values provably in range (e.g. derived from a [`chrono::DateTime`]), or for decode
116    /// paths that must stay infallible and knowingly tolerate legacy out-of-range values.
117    pub fn from_micros_uncheck(timestamp_micros: i64) -> Self {
118        Self(timestamp_micros)
119    }
120
121    /// Creates a `Timestamptz` from nanoseconds, flooring to microseconds.
122    ///
123    /// An `i64` nanosecond count spans only ±292 years around the epoch, well within the
124    /// representable range, so this cannot fail.
125    pub fn from_nanos(timestamp_nanos: i64) -> Self {
126        Self(timestamp_nanos.div_euclid(1_000))
127    }
128
129    /// Returns the number of non-leap-microseconds since January 1, 1970 UTC.
130    pub fn timestamp_micros(&self) -> i64 {
131        self.0
132    }
133
134    /// Returns the number of non-leap-milliseconds since January 1, 1970 UTC.
135    pub fn timestamp_millis(&self) -> i64 {
136        self.0.div_euclid(1_000)
137    }
138
139    /// Returns the number of non-leap-nanosseconds since January 1, 1970 UTC.
140    pub fn timestamp_nanos(&self) -> Option<i64> {
141        self.0.checked_mul(1_000)
142    }
143
144    /// Returns the number of non-leap seconds since January 1, 1970 0:00:00 UTC (aka "UNIX
145    /// timestamp").
146    pub fn timestamp(&self) -> i64 {
147        self.0.div_euclid(1_000_000)
148    }
149
150    /// Returns the number of nanoseconds since the last second boundary.
151    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        // Every fallible constructor validates representability, so this can fail only for
189        // values built by `from_micros_uncheck`: a contract violation, or a legacy out-of-range
190        // value read back from storage (#26397).
191        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        // Try `speedate` first
205        // * It is also used by `str_to_{date,time,timestamp}`
206        // * It can parse without seconds `2006-01-02 15:04-07:00`
207        let ret = match speedate::DateTime::parse_str_rfc3339(s) {
208            Ok(r) => r,
209            Err(_) => {
210                // Supplement with `chrono` for existing cases:
211                // * Extra space before offset `2006-01-02 15:04:05 -07:00`
212                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        // The range check must accept exactly what `DateTime::from_timestamp_micros` accepts.
267        for micros in [
268            0,
269            i64::MAX,
270            i64::MIN,
271            8_210_266_876_799_999_999,  // chrono max
272            8_210_266_876_800_000_000,  // chrono max + 1
273            -8_334_601_228_800_000_000, // chrono min
274            -8_334_601_228_800_000_001, // chrono min - 1
275        ] {
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        // Most standard: ISO 8601 & RFC 3339
294        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        // RFC 3339
298        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        // PostgreSQL, but neither ISO 8601 nor RFC 3339
302        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}