Skip to main content

pg_srv/values/
timestamp.rs

1//! Timestamp value representation for PostgreSQL protocol
2
3use crate::{
4    protocol::{ErrorCode, ErrorResponse},
5    FromProtocolValue, ProtocolError, ToProtocolValue,
6};
7use byteorder::{BigEndian, ByteOrder};
8use bytes::{BufMut, BytesMut};
9use chrono::{
10    format::{
11        Fixed, Item,
12        Numeric::{Day, Hour, Minute, Month, Second, Year},
13        Pad::Zero,
14    },
15    prelude::*,
16};
17use chrono_tz::Tz;
18use std::backtrace::Backtrace;
19use std::io::Error;
20use std::{
21    fmt::{self, Debug, Display, Formatter},
22    io,
23};
24
25#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
26pub struct TimestampValue {
27    unix_nano: i64,
28    tz: Option<String>,
29}
30
31impl TimestampValue {
32    pub fn new(mut unix_nano: i64, tz: Option<String>) -> TimestampValue {
33        // This is a hack to workaround a mismatch between on-disk and in-memory representations.
34        // We use microsecond precision on-disk.
35        unix_nano -= unix_nano % 1000;
36        TimestampValue { unix_nano, tz }
37    }
38
39    pub fn to_naive_datetime(&self) -> NaiveDateTime {
40        // Convert nanoseconds to seconds and nanoseconds
41        let secs = self.unix_nano / 1_000_000_000;
42        let nsecs = (self.unix_nano % 1_000_000_000) as u32;
43        DateTime::from_timestamp(secs, nsecs)
44            .unwrap_or_else(|| panic!("Invalid timestamp: {}", self.unix_nano))
45            .naive_utc()
46    }
47
48    pub fn to_fixed_datetime(&self) -> io::Result<DateTime<Tz>> {
49        assert!(self.tz.is_some());
50        let tz = self
51            .tz
52            .as_ref()
53            .unwrap()
54            .parse::<Tz>()
55            .map_err(|err| io::Error::other(err.to_string()))?;
56        let ndt = self.to_naive_datetime();
57        Ok(tz.from_utc_datetime(&ndt))
58    }
59
60    pub fn tz_ref(&self) -> &Option<String> {
61        &self.tz
62    }
63
64    pub fn get_time_stamp(&self) -> i64 {
65        self.unix_nano
66    }
67}
68
69impl Debug for TimestampValue {
70    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
71        f.debug_struct("TimestampValue")
72            .field("unix_nano", &self.unix_nano)
73            .field("tz", &self.tz)
74            .field("str", &self.to_string())
75            .finish()
76    }
77}
78
79impl Display for TimestampValue {
80    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
81        let formatted = Utc.timestamp_nanos(self.unix_nano).format_with_items(
82            [
83                Item::Numeric(Year, Zero),
84                Item::Literal("-"),
85                Item::Numeric(Month, Zero),
86                Item::Literal("-"),
87                Item::Numeric(Day, Zero),
88                Item::Literal("T"),
89                Item::Numeric(Hour, Zero),
90                Item::Literal(":"),
91                Item::Numeric(Minute, Zero),
92                Item::Literal(":"),
93                Item::Numeric(Second, Zero),
94                Item::Fixed(Fixed::Nanosecond3),
95            ]
96            .iter(),
97        );
98        write!(f, "{}", formatted)
99    }
100}
101
102// POSTGRES_EPOCH_JDATE
103// https://github.com/postgres/postgres/blob/REL_14_4/src/include/datatype/timestamp.h#L163
104pub(crate) fn pg_base_date_epoch() -> NaiveDateTime {
105    NaiveDate::from_ymd_opt(2000, 1, 1)
106        .unwrap()
107        .and_hms_opt(0, 0, 0)
108        .unwrap()
109}
110
111impl ToProtocolValue for TimestampValue {
112    fn to_text(&self, buf: &mut BytesMut) -> Result<(), ProtocolError> {
113        let ndt = match self.tz_ref() {
114            None => self.to_naive_datetime(),
115            Some(_) => self.to_fixed_datetime()?.naive_utc(),
116        };
117
118        // 2022-04-25 15:36:49.39705+00
119        let as_str = ndt.format("%Y-%m-%d %H:%M:%S%.6f").to_string();
120
121        match self.tz_ref() {
122            None => as_str.to_text(buf),
123            Some(_) => (as_str + "+00").to_text(buf),
124        }
125    }
126
127    // https://github.com/postgres/postgres/blob/REL_14_4/src/backend/utils/adt/timestamp.c#L267
128    fn to_binary(&self, buf: &mut BytesMut) -> Result<(), ProtocolError> {
129        let ndt = match self.tz_ref() {
130            None => self.to_naive_datetime(),
131            Some(_) => self.to_fixed_datetime()?.naive_utc(),
132        };
133
134        let n = ndt
135            .signed_duration_since(pg_base_date_epoch())
136            .num_microseconds()
137            .ok_or(Error::other(
138                "Unable to extract number of seconds from timestamp",
139            ))?;
140
141        buf.put_i32(8);
142        buf.put_i64(n);
143
144        Ok(())
145    }
146}
147
148impl FromProtocolValue for TimestampValue {
149    fn from_text(raw: &[u8]) -> Result<Self, ProtocolError> {
150        let as_str = std::str::from_utf8(raw).map_err(|err| ProtocolError::ErrorResponse {
151            source: ErrorResponse::error(ErrorCode::ProtocolViolation, err.to_string()),
152            backtrace: Backtrace::capture(),
153        })?;
154
155        // Parse timestamp string in format "YYYY-MM-DD HH:MM:SS[.fff]", but PostgreSQL supports
156        // more formats, so let's align this with parse_date_str function from cubesql crate.
157        let parsed_datetime = NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S")
158            .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S%.f"))
159            .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S%.f UTC"))
160            .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S"))
161            .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S%.f"))
162            .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S%.fZ"))
163            .or_else(|_| {
164                NaiveDate::parse_from_str(as_str, "%Y-%m-%d").map(|date| {
165                    date.and_hms_opt(0, 0, 0)
166                        .expect("Unable to set time to 00:00:00")
167                })
168            })
169            .map_err(|err| ProtocolError::ErrorResponse {
170                source: ErrorResponse::error(
171                    ErrorCode::ProtocolViolation,
172                    format!(
173                        "Unable to parse timestamp from text: '{}', error: {}",
174                        as_str, err
175                    ),
176                ),
177                backtrace: Backtrace::capture(),
178            })?;
179
180        // Convert to Unix nanoseconds
181        let unix_nano = parsed_datetime
182            .and_utc()
183            .timestamp_nanos_opt()
184            .ok_or_else(|| ProtocolError::ErrorResponse {
185                source: ErrorResponse::error(
186                    ErrorCode::ProtocolViolation,
187                    format!("Timestamp out of range: '{}'", as_str),
188                ),
189                backtrace: Backtrace::capture(),
190            })?;
191
192        Ok(TimestampValue::new(unix_nano, None))
193    }
194
195    // https://github.com/postgres/postgres/blob/REL_14_4/src/backend/utils/adt/timestamp.c#L234
196    fn from_binary(raw: &[u8]) -> Result<Self, ProtocolError> {
197        if raw.len() != 8 {
198            return Err(ProtocolError::ErrorResponse {
199                source: ErrorResponse::error(
200                    ErrorCode::ProtocolViolation,
201                    format!(
202                        "Invalid binary timestamp length: expected 8 bytes, got {}",
203                        raw.len()
204                    ),
205                ),
206                backtrace: Backtrace::capture(),
207            });
208        }
209
210        let pg_microseconds = BigEndian::read_i64(raw);
211
212        // Convert PostgreSQL microseconds to Unix nanoseconds
213        let unix_nano = pg_base_date_epoch()
214            .and_utc()
215            .timestamp_nanos_opt()
216            .expect("Unable to get timestamp nanos for pg_base_date_epoch")
217            + (pg_microseconds * 1_000);
218
219        Ok(TimestampValue::new(unix_nano, None))
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use crate::ProtocolError;
227
228    #[test]
229    fn test_timestamp_creation() -> Result<(), ProtocolError> {
230        let ts = TimestampValue::new(1650890322000000000, None);
231        assert_eq!(ts.get_time_stamp(), 1650890322000000000);
232        assert_eq!(ts.tz_ref(), &None);
233
234        let ts_with_tz = TimestampValue::new(1650890322000000000, Some("UTC".to_string()));
235        assert_eq!(ts_with_tz.get_time_stamp(), 1650890322000000000);
236        assert_eq!(ts_with_tz.tz_ref(), &Some("UTC".to_string()));
237
238        Ok(())
239    }
240
241    #[test]
242    fn test_timestamp_to_string() {
243        let ts = TimestampValue::new(1650890322000000000, None);
244        // The string representation should match the expected format
245        assert!(!ts.to_string().is_empty());
246    }
247
248    #[test]
249    fn test_timestamp_precision_hack() {
250        // Test that nanoseconds are truncated to milliseconds
251        let ts = TimestampValue::new(1650890322123456789, None);
252        assert_eq!(ts.get_time_stamp(), 1650890322123456000);
253    }
254
255    #[test]
256    fn test_invalid_timestamp_text() {
257        // Test that invalid text formats return errors
258        assert!(TimestampValue::from_text(b"invalid-date").is_err());
259        assert!(TimestampValue::from_text(b"2025-13-45 25:70:99").is_err());
260        assert!(TimestampValue::from_text(b"").is_err());
261    }
262
263    #[test]
264    fn test_timestamp_from_text_various_formats() {
265        // Test basic format without fractional seconds
266        let ts1 = TimestampValue::from_text(b"2025-08-04 20:15:47").unwrap();
267        assert_eq!(ts1.to_naive_datetime().to_string(), "2025-08-04 20:15:47");
268
269        // Test PostgreSQL format with 6-digit fractional seconds
270        let ts2 = TimestampValue::from_text(b"2025-08-04 20:16:54.853660").unwrap();
271        assert_eq!(
272            ts2.to_naive_datetime()
273                .format("%Y-%m-%d %H:%M:%S%.6f")
274                .to_string(),
275            "2025-08-04 20:16:54.853660"
276        );
277
278        // Test format with 3 fractional seconds
279        let ts3 = TimestampValue::from_text(b"2025-08-04 20:15:47.953").unwrap();
280        assert_eq!(
281            ts3.to_naive_datetime()
282                .format("%Y-%m-%d %H:%M:%S%.3f")
283                .to_string(),
284            "2025-08-04 20:15:47.953"
285        );
286
287        // Test ISO format with T separator
288        let ts4 = TimestampValue::from_text(b"2025-08-04T20:15:47").unwrap();
289        assert_eq!(ts4.to_naive_datetime().to_string(), "2025-08-04 20:15:47");
290
291        // Test ISO format with T separator and fractional seconds
292        let ts5 = TimestampValue::from_text(b"2025-08-04T20:15:47.953116").unwrap();
293        assert_eq!(
294            ts5.to_naive_datetime()
295                .format("%Y-%m-%d %H:%M:%S%.6f")
296                .to_string(),
297            "2025-08-04 20:15:47.953116"
298        );
299    }
300
301    #[test]
302    fn test_invalid_timestamp_binary() {
303        // Test that invalid binary data returns errors
304        assert!(TimestampValue::from_binary(&[1, 2, 3]).is_err()); // Wrong length
305        assert!(TimestampValue::from_binary(&[]).is_err()); // Empty
306    }
307}