1use 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 unix_nano -= unix_nano % 1000;
36 TimestampValue { unix_nano, tz }
37 }
38
39 pub fn to_naive_datetime(&self) -> NaiveDateTime {
40 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
102pub(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 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 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 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 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 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 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 assert!(!ts.to_string().is_empty());
246 }
247
248 #[test]
249 fn test_timestamp_precision_hack() {
250 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 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 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 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 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 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 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 assert!(TimestampValue::from_binary(&[1, 2, 3]).is_err()); assert!(TimestampValue::from_binary(&[]).is_err()); }
307}