From e2b5784ea1490c3211b157446e49106a9271df8c Mon Sep 17 00:00:00 2001
From: Elvis Pranskevichus <elvis@edgedb.com>
Date: Wed, 13 Nov 2024 22:48:14 -0800
Subject: [PATCH] Implement codec support for newly exposed `std::pg::` types

Non-binary `json` type and raw datetime types which all map to existing
model types, so this is trivial.
---
 edgedb-protocol/src/codec.rs | 37 +++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/edgedb-protocol/src/codec.rs b/edgedb-protocol/src/codec.rs
index 6180e506..7eddbe10 100644
--- a/edgedb-protocol/src/codec.rs
+++ b/edgedb-protocol/src/codec.rs
@@ -11,7 +11,7 @@ use std::str;
 use std::sync::Arc;
 
 use bytes::{Buf, BufMut, BytesMut};
-use snafu::{ensure, OptionExt};
+use snafu::{ensure, OptionExt, ResultExt};
 use uuid::Uuid as UuidVal;
 
 use crate::common::Cardinality;
@@ -44,6 +44,11 @@ pub const STD_JSON: UuidVal = UuidVal::from_u128(0x10f);
 pub const STD_BIGINT: UuidVal = UuidVal::from_u128(0x110);
 pub const CFG_MEMORY: UuidVal = UuidVal::from_u128(0x130);
 pub const PGVECTOR_VECTOR: UuidVal = UuidVal::from_u128(0x9565dd88_04f5_11ee_a691_0b6ebe179825);
+pub const STD_PG_JSON: UuidVal = UuidVal::from_u128(0x1000001);
+pub const STD_PG_TIMESTAMPTZ: UuidVal = UuidVal::from_u128(0x1000002);
+pub const STD_PG_TIMESTAMP: UuidVal = UuidVal::from_u128(0x1000003);
+pub const STD_PG_DATE: UuidVal = UuidVal::from_u128(0x1000004);
+pub const STD_PG_INTERVAL: UuidVal = UuidVal::from_u128(0x1000005);
 
 pub(crate) fn uuid_to_known_name(uuid: &UuidVal) -> Option<&'static str> {
     match *uuid {
@@ -68,6 +73,11 @@ pub(crate) fn uuid_to_known_name(uuid: &UuidVal) -> Option<&'static str> {
         STD_BIGINT => Some("BaseScalar(bigint)"),
         CFG_MEMORY => Some("BaseScalar(cfg::memory)"),
         PGVECTOR_VECTOR => Some("BaseScalar(ext::pgvector::vector)"),
+        STD_PG_JSON => Some("BaseScalar(std::pg::json)"),
+        STD_PG_TIMESTAMPTZ => Some("BaseScalar(std::pg::timestamptz)"),
+        STD_PG_TIMESTAMP => Some("BaseScalar(std::pg::timestamp)"),
+        STD_PG_DATE => Some("BaseScalar(std::pg::date)"),
+        STD_PG_INTERVAL => Some("BaseScalar(std::pg::interval)"),
         _ => None,
     }
 }
@@ -182,6 +192,9 @@ pub struct Bool;
 #[derive(Debug)]
 pub struct Json;
 
+#[derive(Debug)]
+pub struct PgTextJson;
+
 #[derive(Debug)]
 pub struct Nothing;
 
@@ -358,6 +371,11 @@ pub fn scalar_codec(uuid: &UuidVal) -> Result<Arc<dyn Codec>, CodecError> {
         STD_BIGINT => Ok(Arc::new(BigInt {})),
         CFG_MEMORY => Ok(Arc::new(ConfigMemory {})),
         PGVECTOR_VECTOR => Ok(Arc::new(Vector {})),
+        STD_PG_JSON => Ok(Arc::new(PgTextJson {})),
+        STD_PG_TIMESTAMPTZ => Ok(Arc::new(Datetime {})),
+        STD_PG_TIMESTAMP => Ok(Arc::new(LocalDatetime {})),
+        STD_PG_DATE => Ok(Arc::new(LocalDate {})),
+        STD_PG_INTERVAL => Ok(Arc::new(RelativeDuration {})),
         _ => errors::UndefinedBaseScalar { uuid: uuid.clone() }.fail()?,
     }
 }
@@ -1145,6 +1163,23 @@ impl Codec for Json {
     }
 }
 
+impl Codec for PgTextJson {
+    fn decode(&self, buf: &[u8]) -> Result<Value, DecodeError> {
+        let val = str::from_utf8(buf).context(errors::InvalidUtf8)?.to_owned();
+        let json = model::Json::new_unchecked(val);
+        Ok(Value::Json(json))
+    }
+    fn encode(&self, buf: &mut BytesMut, val: &Value) -> Result<(), EncodeError> {
+        let val = match val {
+            Value::Json(val) => val,
+            _ => Err(errors::invalid_value(type_name::<Self>(), val))?,
+        };
+        buf.reserve(val.len());
+        buf.extend(val.as_bytes());
+        Ok(())
+    }
+}
+
 impl Codec for Scalar {
     fn decode(&self, buf: &[u8]) -> Result<Value, DecodeError> {
         self.inner.decode(buf)