feat: compatible with postgres interval type

src/servers/src/postgres/types.rs

@@ -13,9 +13,12 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::fmt::Display;
 use std::ops::Deref;
 
+use bytes::{Buf, BufMut};
 use chrono::{NaiveDate, NaiveDateTime};
+use common_time::Interval;
 use datafusion_common::ScalarValue;
 use datatypes::prelude::{ConcreteDataType, Value};
 use datatypes::schema::Schema;
@@ -24,6 +27,8 @@ use pgwire::api::portal::{Format, Portal};
 use pgwire::api::results::{DataRowEncoder, FieldInfo};
 use pgwire::api::Type;
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
+use pgwire::types::ToSqlText;
+use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql};
 use query::plan::LogicalPlan;
 
 use crate::error::{self, Error, Result};
@@ -98,14 +103,13 @@ pub(super) fn encode_value(value: &Value, builder: &mut DataRowEncoder) -> PgWir
                 })))
             }
         }
-        Value::Interval(_) | Value::List(_) => {
-            Err(PgWireError::ApiError(Box::new(Error::Internal {
-                err_msg: format!(
-                    "cannot write value {:?} in postgres protocol: unimplemented",
-                    &value
-                ),
-            })))
-        }
+        Value::Interval(v) => builder.encode_field(&PgInterval::from(*v)),
+        Value::List(_) => Err(PgWireError::ApiError(Box::new(Error::Internal {
+            err_msg: format!(
+                "cannot write value {:?} in postgres protocol: unimplemented",
+                &value
+            ),
+        }))),
     }
 }
 
@@ -195,6 +199,10 @@ pub(super) fn parameter_to_string(portal: &Portal<SqlPlan>, idx: usize) -> PgWir
             .parameter::<NaiveDateTime>(idx, param_type)?
             .map(|v| v.format("%Y-%m-%d %H:%M:%S%.6f").to_string())
             .unwrap_or_else(|| "".to_owned())),
+        &Type::INTERVAL => Ok(portal
+            .parameter::<PgInterval>(idx, param_type)?
+            .map(|v| v.to_string())
+            .unwrap_or_else(|| "".to_owned())),
         _ => Err(invalid_parameter_error(
             "unsupported_parameter_type",
             Some(&param_type.to_string()),
@@ -478,6 +486,23 @@ pub(super) fn parameters_to_scalar_values(
                     }
                 }
             }
+            &Type::INTERVAL => {
+                let data = portal.parameter::<PgInterval>(idx, &client_type)?;
+                match server_type {
+                    ConcreteDataType::Interval(_) => {
+                        ScalarValue::IntervalMonthDayNano(data.map(|i| Interval::from(i).to_i128()))
+                    }
+                    _ => {
+                        return Err(invalid_parameter_error(
+                            "invalid_parameter_type",
+                            Some(&format!(
+                                "Expected: {}, found: {}",
+                                server_type, client_type
+                            )),
+                        ));
+                    }
+                }
+            }
             &Type::BYTEA => {
                 let data = portal.parameter::<Vec<u8>>(idx, &client_type)?;
                 match server_type {
@@ -524,6 +549,113 @@ pub(super) fn param_types_to_pg_types(
     Ok(types)
 }
 
+#[derive(Debug, Clone, Copy, Default)]
+pub struct PgInterval {
+    months: i32,
+    days: i32,
+    microseconds: i64,
+}
+
+impl From<Interval> for PgInterval {
+    fn from(interval: Interval) -> Self {
+        let (months, days, nanos) = interval.to_month_day_nano();
+        Self {
+            months,
+            days,
+            microseconds: nanos / 1000,
+        }
+    }
+}
+
+impl From<PgInterval> for Interval {
+    fn from(interval: PgInterval) -> Self {
+        Interval::from_month_day_nano(
+            interval.months,
+            interval.days,
+            // Maybe overflow, but most scenarios ok.
+            interval.microseconds.checked_mul(1000).unwrap_or_else(|| {
+                if interval.microseconds.is_negative() {
+                    i64::MIN
+                } else {
+                    i64::MAX
+                }
+            }),
+        )
+    }
+}
+
+impl Display for PgInterval {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", Interval::from(*self).to_postgres_string())
+    }
+}
+
+impl ToSql for PgInterval {
+    to_sql_checked!();
+
+    fn to_sql(
+        &self,
+        _: &Type,
+        out: &mut bytes::BytesMut,
+    ) -> std::result::Result<postgres_types::IsNull, Box<dyn snafu::Error + Sync + Send>>
+    where
+        Self: Sized,
+    {
+        // https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/timestamp.c#L989-L991
+        out.put_i64(self.microseconds);
+        out.put_i32(self.days);
+        out.put_i32(self.months);
+        Ok(postgres_types::IsNull::No)
+    }
+
+    fn accepts(ty: &Type) -> bool
+    where
+        Self: Sized,
+    {
+        matches!(ty, &Type::INTERVAL)
+    }
+}
+
+impl<'a> FromSql<'a> for PgInterval {
+    fn from_sql(
+        _: &Type,
+        mut raw: &'a [u8],
+    ) -> std::result::Result<Self, Box<dyn snafu::Error + Sync + Send>> {
+        // https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/timestamp.c#L1007-L1010
+        let microseconds = raw.get_i64();
+        let days = raw.get_i32();
+        let months = raw.get_i32();
+        Ok(PgInterval {
+            months,
+            days,
+            microseconds,
+        })
+    }
+
+    fn accepts(ty: &Type) -> bool {
+        matches!(ty, &Type::INTERVAL)
+    }
+}
+
+impl ToSqlText for PgInterval {
+    fn to_sql_text(
+        &self,
+        ty: &Type,
+        out: &mut bytes::BytesMut,
+    ) -> std::result::Result<postgres_types::IsNull, Box<dyn snafu::Error + Sync + Send>>
+    where
+        Self: Sized,
+    {
+        let fmt = match ty {
+            &Type::INTERVAL => self.to_string(),
+            _ => return Err("unsupported type".into()),
+        };
+
+        out.put_slice(fmt.as_bytes());
+        Ok(IsNull::No)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
@@ -559,6 +691,11 @@ mod test {
             ),
             ColumnSchema::new("dates", ConcreteDataType::date_datatype(), true),
             ColumnSchema::new("times", ConcreteDataType::time_second_datatype(), true),
+            ColumnSchema::new(
+                "intervals",
+                ConcreteDataType::interval_month_day_nano_datatype(),
+                true,
+            ),
         ];
         let pg_field_info = vec![
             FieldInfo::new("nulls".into(), None, None, Type::UNKNOWN, FieldFormat::Text),
@@ -608,6 +745,13 @@ mod test {
             ),
             FieldInfo::new("dates".into(), None, None, Type::DATE, FieldFormat::Text),
             FieldInfo::new("times".into(), None, None, Type::TIME, FieldFormat::Text),
+            FieldInfo::new(
+                "intervals".into(),
+                None,
+                None,
+                Type::INTERVAL,
+                FieldFormat::Text,
+            ),
         ];
         let schema = Schema::new(column_schemas);
         let fs = schema_to_pg(&schema, &Format::UnifiedText).unwrap();
@@ -703,6 +847,13 @@ mod test {
                 Type::TIMESTAMP,
                 FieldFormat::Text,
             ),
+            FieldInfo::new(
+                "intervals".into(),
+                None,
+                None,
+                Type::INTERVAL,
+                FieldFormat::Text,
+            ),
         ];
 
         let values = vec![
@@ -732,6 +883,7 @@ mod test {
             Value::Time(1001i64.into()),
             Value::DateTime(1000001i64.into()),
             Value::Timestamp(1000001i64.into()),
+            Value::Interval(1000001i128.into()),
         ];
         let mut builder = DataRowEncoder::new(Arc::new(schema));
         for i in values.iter() {