fix(clickhouse): timestamptz type formatting issue (supabase#535)

Author: burmecia

wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs

@@ -28,6 +28,17 @@ use supabase_wrappers::prelude::*;
 
 use super::{ClickHouseFdwError, ClickHouseFdwResult};
 
+struct Formatter;
+
+impl CellFormatter for Formatter {
+    fn fmt_cell(&mut self, cell: &Cell) -> String {
+        match cell {
+            Cell::Timestamptz(tstz) => format!("parseDateTime64BestEffort('{tstz}', 6)"),
+            _ => format!("{cell}"),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 struct ConvertedRow {
     values: Vec<Option<Cell>>,
@@ -138,7 +149,7 @@ fn convert_row_simple(
                     // convert from Unix timestamp (seconds since 1970-01-01) to
                     // PostgreSQL timestamp (microseconds since 2000-01-01)
                     // difference between 1970 and 2000 epoch: 946684800 seconds
-                    let pg_raw_timestamp = (value.timestamp() - 946684800) * 1_000_000;
+                    let pg_raw_timestamp = value.timestamp_micros() - 946684800 * 1_000_000;
                     let pg_timestamp = pgrx::datum::Timestamp::saturating_from_raw(
                         pg_raw_timestamp as pg_sys::Timestamp,
                     );
@@ -445,6 +456,7 @@ impl ClickHouseFdw {
         let mut sql = format!("select {} from {}", tgts, &table);
 
         if !quals.is_empty() {
+            let mut formatter = Formatter {};
             let cond = quals
                 .iter()
                 .filter(|q| {
@@ -455,7 +467,7 @@ impl ClickHouseFdw {
                     };
                     !is_param && !is_array
                 })
-                .map(|q| q.deparse())
+                .map(|q| q.deparse_with_fmt(&mut formatter))
                 .collect::<Vec<String>>()
                 .join(" and ");
 
@@ -850,6 +862,17 @@ impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
                         };
                         Ok(val)
                     }
+                    Cell::Timestamptz(_) => {
+                        let s = c.to_string().replace('\'', "");
+                        let tm = DateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S%.6f%#z")?;
+                        let utc_tm = tm.with_timezone(&Utc);
+                        let val = if is_nullable {
+                            ChValue::Nullable(either::Either::Right(Box::new(utc_tm.into())))
+                        } else {
+                            ChValue::from(utc_tm)
+                        };
+                        Ok(val)
+                    }
                     Cell::Uuid(v) => {
                         let uuid = Uuid::try_parse(&v.to_string())?;
                         let val = if is_nullable {
@@ -905,13 +928,15 @@ impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
         let pool = Pool::new(self.conn_str.clone());
         let mut client = self.rt.block_on(pool.get_handle())?;
 
+        let mut formatter = Formatter {};
         let mut sets = Vec::new();
         for (col, cell) in new_row.iter() {
             if col == &self.rowid_col {
                 continue;
             }
             if let Some(cell) = cell {
-                sets.push(format!("{col} = {cell}"));
+                let formatted = formatter.fmt_cell(cell);
+                sets.push(format!("{col} = {formatted}"));
             } else {
                 sets.push(format!("{col} = null"));
             }

wrappers/src/fdw/clickhouse_fdw/tests.rs

@@ -3,7 +3,10 @@
 mod tests {
     use clickhouse_rs as ch;
     use pgrx::prelude::*;
-    use pgrx::{datum::Timestamp, pg_test, Uuid};
+    use pgrx::{
+        datum::{Timestamp, TimestampWithTimeZone},
+        pg_test, Uuid,
+    };
     use supabase_wrappers::prelude::create_async_runtime;
 
     #[pg_test]
@@ -37,7 +40,8 @@ mod tests {
                             u16col Nullable(UInt16),
                             i32col Nullable(Int32),
                             u32col Nullable(UInt32),
-                            created_at DateTime('UTC')
+                            created_at DateTime('UTC'),
+                            updated_at DateTime64(6, 'Asia/Singapore')
                         ) engine = Memory",
                     )
                     .await
@@ -80,7 +84,8 @@ mod tests {
                     u16col integer,
                     i32col integer,
                     u32col bigint,
-                    created_at timestamp
+                    created_at timestamp,
+                    updated_at timestamptz
                   )
                   SERVER my_clickhouse_server
                   OPTIONS (
@@ -112,7 +117,8 @@ mod tests {
                     u16col integer,
                     i32col integer,
                     u32col bigint,
-                    created_at timestamp
+                    created_at timestamp,
+                    updated_at timestamptz
                   )
                   SERVER my_clickhouse_server
                   OPTIONS (
@@ -144,7 +150,8 @@ mod tests {
                     u16col integer,
                     i32col integer,
                     u32col bigint,
-                    created_at timestamp
+                    created_at timestamp,
+                    updated_at timestamptz
                   )
                   SERVER my_clickhouse_server
                   OPTIONS (
@@ -190,9 +197,9 @@ mod tests {
             c.update(
                 "INSERT INTO test_table (id, name, amt, uid, fstr, bignum, dnum,
                     arr_i64, arr_str, is_valid, i8col, u8col, i16col, u16col,
-                    i32col, u32col, created_at)
+                    i32col, u32col, created_at, updated_at)
                  VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
-                     $14, $15, $16, $17)",
+                     $14, $15, $16, $17, $18)",
                 None,
                 &[
                     42.into(),
@@ -212,6 +219,9 @@ mod tests {
                     46i32.into(),
                     47i64.into(),
                     Timestamp::new(2025, 5, 1, 2, 3, 4.0).into(),
+                    TimestampWithTimeZone::with_timezone(2025, 5, 1, 2, 3, 4.0, "Asia/Singapore")
+                        .unwrap()
+                        .into(),
                 ],
             )
             .unwrap();
@@ -443,6 +453,29 @@ mod tests {
                 Some(Uuid::from_bytes([43u8; 16])),
             );
 
+            // test timestamptz data type
+            c.update(
+                "UPDATE test_table
+                 SET updated_at = '2025-05-01 01:02:03.112233+8'::timestamptz
+                 WHERE id = 42
+                ",
+                None,
+                &[],
+            )
+            .unwrap();
+            assert_eq!(
+                c.select(
+                    "SELECT id FROM test_table WHERE updated_at = '2025-05-01 01:02:03.112233+8'::timestamptz",
+                    None,
+                    &[]
+                )
+                .unwrap()
+                .first()
+                .get_one::<i64>()
+                .unwrap(),
+                Some(42),
+            );
+
             // test delete data in foreign table
             c.update("DELETE FROM test_table WHERE id = 42", None, &[])
                 .unwrap();