From 2997018005edecae8597eff3433acf049a5cdfb7 Mon Sep 17 00:00:00 2001
From: George MacKerron <george@mackerron.co.uk>
Date: Mon, 8 May 2023 11:02:12 +0100
Subject: [PATCH] Added support to return timestamptz, timestamp, date and time
 for http queries

---
 Cargo.lock                   |  1 +
 Cargo.toml                   |  2 +-
 proxy/src/http/pg_to_json.rs | 16 ++++++++++++++++
 proxy/src/http/websocket.rs  | 11 +++++------
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9f77612b31..c148fb3d0b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2886,6 +2886,7 @@ version = "0.2.4"
 source = "git+https://github.com/neondatabase/rust-postgres.git?rev=43e6db254a97fdecbce33d8bc0890accfd74495e#43e6db254a97fdecbce33d8bc0890accfd74495e"
 dependencies = [
  "bytes",
+ "chrono",
  "fallible-iterator",
  "postgres-protocol",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 9811bc1d2d..520e48baad 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -124,7 +124,7 @@ env_logger = "0.10"
 log = "0.4"
 
 ## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e", features = ["with-chrono-0_4"] }
 postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
 postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
 tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
diff --git a/proxy/src/http/pg_to_json.rs b/proxy/src/http/pg_to_json.rs
index b9b73f9c3a..815cc17c64 100644
--- a/proxy/src/http/pg_to_json.rs
+++ b/proxy/src/http/pg_to_json.rs
@@ -1,5 +1,6 @@
 
 use anyhow::{anyhow, Context};
+use chrono::{Utc, DateTime, NaiveDateTime, NaiveTime, NaiveDate};
 use serde_json::Map;
 use tokio_postgres::{
     types::{FromSql, Type},
@@ -39,6 +40,21 @@ pub fn pg_cell_to_json_value(
         // for rust-postgres <> postgres type-mappings: https://docs.rs/postgres/latest/postgres/types/trait.FromSql.html#types
         // for postgres types: https://www.postgresql.org/docs/7.4/datatype.html#DATATYPE-TABLE
 
+        Type::TIMESTAMPTZ => get_basic(row, column, column_i, |a: DateTime<Utc>| {
+            Ok(JSONValue::String(a.to_rfc3339()))
+        })?,
+        Type::TIMESTAMP => get_basic(row, column, column_i, |a: NaiveDateTime| {
+            Ok(JSONValue::String(a.format("%Y-%m-%dT%H:%M:%S%.6f").to_string()))
+        })?,
+        
+        Type::TIME => get_basic(row, column, column_i, |a: NaiveTime| {
+            Ok(JSONValue::String(a.format("%H:%M:%S%.6f").to_string()))
+        })?,
+        Type::DATE => get_basic(row, column, column_i, |a: NaiveDate| {
+            Ok(JSONValue::String(a.format("%Y-%m-%d").to_string()))
+        })?,
+        // no TIMETZ support?
+
         // single types
         Type::BOOL => get_basic(row, column, column_i, |a: bool| Ok(JSONValue::Bool(a)))?,
         Type::INT2 => get_basic(row, column, column_i, |a: i16| {
diff --git a/proxy/src/http/websocket.rs b/proxy/src/http/websocket.rs
index c10d1b8bfa..3d9850fed9 100644
--- a/proxy/src/http/websocket.rs
+++ b/proxy/src/http/websocket.rs
@@ -273,10 +273,9 @@ async fn handle_sql(
         query: String,
         params: Vec<serde_json::Value>
     }
-
     let query_data: QueryData = serde_json::from_slice(&data)?;
 
-    let params = StartupMessageParams::new([
+    let credential_params = StartupMessageParams::new([
         ("user", username),
         ("database", dbname),
         ("application_name", "proxy_http_sql"),
@@ -286,7 +285,7 @@ async fn handle_sql(
     let creds = config
         .auth_backend
         .as_ref()
-        .map(|_| auth::ClientCredentials::parse(&params, Some(hostname), common_names))
+        .map(|_| auth::ClientCredentials::parse(&credential_params, Some(hostname), common_names))
         .transpose()?;
 
     let extra = console::ConsoleReqExtra {
@@ -321,7 +320,7 @@ async fn handle_sql(
     });
 
     let query = &query_data.query;
-    let params = query_data.params.iter().map(|value| {
+    let query_params = query_data.params.iter().map(|value| {
         let boxed: Box<dyn ToSql + Sync + Send> = match value {
             Value::Null => Box::new(None::<bool>),
             Value::Bool(b) => Box::new(b.clone()),
@@ -330,10 +329,10 @@ async fn handle_sql(
             _ => panic!("wrong parameter type")
         };
         boxed
-    }).collect::<Vec<Box<dyn ToSql + Sync + Send>>>();
+    }).collect::<Vec<_>>();
 
     let pg_rows: Vec<Row> = client
-        .query_raw(query, params)
+        .query_raw(query, query_params)
         .await?
         .try_collect::<Vec<Row>>()
         .await?;