feat: add timestamp with time zone cell type support

docs/bigquery.md

@@ -15,6 +15,7 @@ The BigQuery Wrapper allows you to read and write data from BigQuery within your
 | date             | DATE          |
 | timestamp        | DATETIME      |
 | timestamp        | TIMESTAMP     |
+| timestamptz      | TIMESTAMP     |
 
 ## Preparation
 

docs/mssql.md

@@ -17,6 +17,7 @@ The SQL Server Wrapper allows you to read data from Microsoft SQL Server within
 | text               | varchar/char/text                |
 | date               | date                             |
 | timestamp          | datetime/datetime2/smalldatetime |
+| timestamptz        | datetime/datetime2/smalldatetime |
 
 ## Preparation
 

docs/s3.md

@@ -34,6 +34,7 @@ The S3 Wrapper uses Parquet file data types from [arrow_array::types](https://do
 | text             | ByteArrayType           |
 | date             | Date64Type              |
 | timestamp        | TimestampNanosecondType |
+| timestamptz      | TimestampNanosecondType |
 
 ## Preparation
 

supabase-wrappers/src/interface.rs

@@ -3,7 +3,7 @@
 
 use crate::FdwRoutine;
 use pgrx::pg_sys::panic::ErrorReport;
-use pgrx::prelude::{Date, Timestamp};
+use pgrx::prelude::{Date, Timestamp, TimestampWithTimeZone};
 use pgrx::{
     fcinfo,
     pg_sys::{self, BuiltinOid, Datum, Oid},
@@ -44,6 +44,7 @@ pub enum Cell {
     String(String),
     Date(Date),
     Timestamp(Timestamp),
+    Timestamptz(TimestampWithTimeZone),
     Json(JsonB),
 }
 
@@ -61,6 +62,7 @@ impl Clone for Cell {
             Cell::String(v) => Cell::String(v.clone()),
             Cell::Date(v) => Cell::Date(*v),
             Cell::Timestamp(v) => Cell::Timestamp(*v),
+            Cell::Timestamptz(v) => Cell::Timestamptz(*v),
             Cell::Json(v) => Cell::Json(JsonB(v.0.clone())),
         }
     }
@@ -94,6 +96,15 @@ impl fmt::Display for Cell {
                 let ts_cstr = CStr::from_ptr(ts.cast_mut_ptr());
                 write!(f, "'{}'", ts_cstr.to_str().unwrap())
             },
+            Cell::Timestamptz(v) => unsafe {
+                let ts = fcinfo::direct_function_call_as_datum(
+                    pg_sys::timestamptz_out,
+                    &[(*v).into_datum()],
+                )
+                .unwrap();
+                let ts_cstr = CStr::from_ptr(ts.cast_mut_ptr());
+                write!(f, "'{}'", ts_cstr.to_str().unwrap())
+            },
             Cell::Json(v) => write!(f, "{:?}", v),
         }
     }
@@ -113,6 +124,7 @@ impl IntoDatum for Cell {
             Cell::String(v) => v.into_datum(),
             Cell::Date(v) => v.into_datum(),
             Cell::Timestamp(v) => v.into_datum(),
+            Cell::Timestamptz(v) => v.into_datum(),
             Cell::Json(v) => v.into_datum(),
         }
     }
@@ -134,6 +146,7 @@ impl IntoDatum for Cell {
             || other == pg_sys::TEXTOID
             || other == pg_sys::DATEOID
             || other == pg_sys::TIMESTAMPOID
+            || other == pg_sys::TIMESTAMPTZOID
             || other == pg_sys::JSONBOID
     }
 }
@@ -181,6 +194,9 @@ impl FromDatum for Cell {
             PgOid::BuiltIn(PgBuiltInOids::TIMESTAMPOID) => Some(Cell::Timestamp(
                 Timestamp::from_datum(datum, false).unwrap(),
             )),
+            PgOid::BuiltIn(PgBuiltInOids::TIMESTAMPTZOID) => Some(Cell::Timestamptz(
+                TimestampWithTimeZone::from_datum(datum, false).unwrap(),
+            )),
             PgOid::BuiltIn(PgBuiltInOids::JSONBOID) => {
                 Some(Cell::Json(JsonB::from_datum(datum, false).unwrap()))
             }

wrappers/src/fdw/airtable_fdw/result.rs

@@ -213,6 +213,19 @@ impl AirtableRecord {
                         }
                     },
                 ),
+                pg_sys::TIMESTAMPTZOID => self.fields.0.get(&col.name).map_or_else(
+                    || Ok(None),
+                    |val| {
+                        if let Value::String(v) = val {
+                            let n = pgrx::TimestampWithTimeZone::from_str(v.as_str())
+                                .ok()
+                                .map(Cell::Timestamptz);
+                            Ok(n)
+                        } else {
+                            Err(())
+                        }
+                    },
+                ),
                 // TODO: Think about adding support for BOOLARRAYOID, NUMERICARRAYOID, TEXTARRAYOID and rest of array types.
                 pg_sys::JSONBOID => self.fields.0.get(&col.name).map_or_else(
                     || Ok(None),

wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs

@@ -393,6 +393,7 @@ impl ForeignDataWrapper<BigQueryFdwError> for BigQueryFdw {
                         Cell::String(v) => row_json[col_name] = json!(v),
                         Cell::Date(v) => row_json[col_name] = json!(v),
                         Cell::Timestamp(v) => row_json[col_name] = json!(v),
+                        Cell::Timestamptz(v) => row_json[col_name] = json!(v),
                         Cell::Json(v) => row_json[col_name] = json!(v),
                     }
                 }

wrappers/src/fdw/logflare_fdw/logflare_fdw.rs

@@ -1,7 +1,7 @@
 use crate::stats;
 use pgrx::{
     pg_sys,
-    prelude::{AnyNumeric, Date, Timestamp},
+    prelude::{AnyNumeric, Date, Timestamp, TimestampWithTimeZone},
 };
 use reqwest::{
     self,
@@ -75,6 +75,10 @@ fn json_value_to_cell(tgt_col: &Column, v: &JsonValue) -> LogflareFdwResult<Cell
             .as_str()
             .and_then(|s| Timestamp::from_str(s).ok())
             .map(Cell::Timestamp),
+        pg_sys::TIMESTAMPTZOID => v
+            .as_str()
+            .and_then(|s| TimestampWithTimeZone::from_str(s).ok())
+            .map(Cell::Timestamptz),
         _ => {
             return Err(LogflareFdwError::UnsupportedColumnType(
                 tgt_col.name.clone(),
@@ -118,6 +122,9 @@ impl LogflareFdw {
                         Cell::String(s) => s.clone(),
                         Cell::Date(d) => d.to_string().as_str().trim_matches('\'').to_owned(),
                         Cell::Timestamp(t) => t.to_string().as_str().trim_matches('\'').to_owned(),
+                        Cell::Timestamptz(t) => {
+                            t.to_string().as_str().trim_matches('\'').to_owned()
+                        }
                         _ => cell.to_string(),
                     };
                     url.query_pairs_mut().append_pair(param_name, &value);

wrappers/src/fdw/mssql_fdw/mssql_fdw.rs

@@ -66,6 +66,12 @@ fn field_to_cell(src_row: &tiberius::Row, tgt_col: &Column) -> MssqlFdwResult<Op
                 Cell::Timestamp(ts.to_utc())
             })
         }
+        PgOid::BuiltIn(PgBuiltInOids::TIMESTAMPTZOID) => {
+            src_row.try_get::<NaiveDateTime, &str>(col_name)?.map(|v| {
+                let ts = to_timestamp(v.timestamp() as f64);
+                Cell::Timestamptz(ts)
+            })
+        }
         _ => {
             return Err(MssqlFdwError::UnsupportedColumnType(tgt_col.name.clone()));
         }

wrappers/src/fdw/s3_fdw/parquet.rs

@@ -339,6 +339,20 @@ impl S3Parquet {
                             })
                         }
                     }
+                    pg_sys::TIMESTAMPTZOID => {
+                        let arr = col
+                            .as_any()
+                            .downcast_ref::<array::TimestampNanosecondArray>()
+                            .ok_or(S3FdwError::ColumnTypeNotMatch(tgt_col.name.clone()))?;
+                        if arr.is_null(self.batch_idx) {
+                            None
+                        } else {
+                            arr.value_as_datetime(self.batch_idx).map(|ts| {
+                                let ts = to_timestamp(ts.timestamp() as f64);
+                                Cell::Timestamptz(ts)
+                            })
+                        }
+                    }
                     _ => return Err(S3FdwError::UnsupportedColumnType(tgt_col.name.clone())),
                 };
                 row.push(&tgt_col.name, cell);