fix: handle native types for joined queries

prisma · Dec 8, 2023 · c667de4 · c667de4
1 parent 1964a5c
commit c667de4
Show file tree

Hide file tree

Showing 12 changed files with 526 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/psl/builtin-connectors/Cargo.toml b/psl/builtin-connectors/Cargo.toml
@@ -13,3 +13,5 @@ indoc.workspace = true
 lsp-types = "0.91.1"
 once_cell = "1.3"
 regex = "1"
+chrono = { version = "0.4.6", default_features = false }
+bigdecimal = "0.3"
diff --git a/psl/builtin-connectors/src/cockroach_datamodel_connector.rs b/psl/builtin-connectors/src/cockroach_datamodel_connector.rs
@@ -1,8 +1,10 @@
 mod native_types;
 mod validations;
 
+use bigdecimal::{BigDecimal, ParseBigDecimalError};
 pub use native_types::CockroachType;
 
+use chrono::*;
 use enumflags2::BitFlags;
 use lsp_types::{CompletionItem, CompletionItemKind, CompletionList};
 use psl_core::{
@@ -307,6 +309,49 @@ impl Connector for CockroachDatamodelConnector {
     fn flavour(&self) -> Flavour {
         Flavour::Cockroach
     }
+
+    fn coerce_json_datetime(
+        &self,
+        str: &str,
+        nt: Option<NativeTypeInstance>,
+    ) -> chrono::ParseResult<chrono::DateTime<FixedOffset>> {
+        let native_type: Option<&CockroachType> = nt.as_ref().map(|nt| nt.downcast_ref());
+
+        match native_type {
+            Some(ct) => match ct {
+                CockroachType::Timestamptz(_) => crate::utils::parse_timestamptz(str),
+                CockroachType::Timestamp(_) => crate::utils::parse_timestamp(str),
+                CockroachType::Date => crate::utils::parse_date(str),
+                CockroachType::Time(_) => crate::utils::parse_time(str),
+                CockroachType::Timetz(_) => {
+                    // We currently don't support time with timezone.
+                    // We strip the timezone information and parse it as a time.
+                    // This is inline with what Quaint does already.
+                    let time_without_tz = str.split("+").next().unwrap();
+
+                    crate::utils::parse_time(time_without_tz)
+                }
+                _ => unreachable!(),
+            },
+            None => crate::utils::parse_timestamptz(str),
+        }
+    }
+
+    fn coerce_json_decimal(
+        &self,
+        str: &str,
+        nt: Option<NativeTypeInstance>,
+    ) -> Result<BigDecimal, ParseBigDecimalError> {
+        let native_type: Option<&CockroachType> = nt.as_ref().map(|nt| nt.downcast_ref());
+
+        match native_type {
+            Some(pt) => match pt {
+                CockroachType::Decimal(_) => crate::utils::parse_decimal(str),
+                _ => unreachable!(),
+            },
+            None => crate::utils::parse_decimal(str),
+        }
+    }
 }
 
 /// An `@default(sequence())` function.

diff --git a/psl/builtin-connectors/src/lib.rs b/psl/builtin-connectors/src/lib.rs
@@ -16,6 +16,7 @@ mod mysql_datamodel_connector;
 mod native_type_definition;
 mod postgres_datamodel_connector;
 mod sqlite_datamodel_connector;
+mod utils;
 
 use psl_core::{datamodel_connector::Connector, ConnectorRegistry};
 

diff --git a/psl/builtin-connectors/src/postgres_datamodel_connector.rs b/psl/builtin-connectors/src/postgres_datamodel_connector.rs
@@ -4,6 +4,8 @@ mod validations;
 
 pub use native_types::PostgresType;
 
+use bigdecimal::{BigDecimal, ParseBigDecimalError};
+use chrono::*;
 use enumflags2::BitFlags;
 use lsp_types::{CompletionItem, CompletionItemKind, CompletionList, InsertTextFormat};
 use psl_core::{
@@ -567,6 +569,47 @@ impl Connector for PostgresDatamodelConnector {
     fn flavour(&self) -> Flavour {
         Flavour::Postgres
     }
+
+    fn coerce_json_datetime(
+        &self,
+        str: &str,
+        nt: Option<NativeTypeInstance>,
+    ) -> chrono::ParseResult<chrono::DateTime<FixedOffset>> {
+        let native_type: Option<&PostgresType> = nt.as_ref().map(|nt| nt.downcast_ref());
+
+        match native_type {
+            Some(pt) => match pt {
+                Timestamptz(_) => crate::utils::parse_timestamptz(str),
+                Timestamp(_) => crate::utils::parse_timestamp(str),
+                Date => crate::utils::parse_date(str),
+                Time(_) => crate::utils::parse_time(str),
+                Timetz(_) => {
+                    let time_without_tz = str.split('+').next().unwrap();
+
+                    crate::utils::parse_time(time_without_tz)
+                }
+                _ => unreachable!(),
+            },
+            None => crate::utils::parse_timestamptz(str),
+        }
+    }
+
+    fn coerce_json_decimal(
+        &self,
+        str: &str,
+        nt: Option<NativeTypeInstance>,
+    ) -> Result<BigDecimal, ParseBigDecimalError> {
+        let native_type: Option<&PostgresType> = nt.as_ref().map(|nt| nt.downcast_ref());
+
+        match native_type {
+            Some(pt) => match pt {
+                Decimal(_) => crate::utils::parse_decimal(str),
+                Money => crate::utils::parse_money(str),
+                _ => unreachable!(),
+            },
+            None => crate::utils::parse_decimal(str),
+        }
+    }
 }
 
 fn allowed_index_operator_classes(algo: IndexAlgorithm, field: walkers::ScalarFieldWalker<'_>) -> Vec<OperatorClass> {

diff --git a/psl/builtin-connectors/src/utils.rs b/psl/builtin-connectors/src/utils.rs
@@ -0,0 +1,39 @@
+use bigdecimal::{BigDecimal, ParseBigDecimalError};
+use chrono::*;
+use std::str::FromStr;
+
+pub(crate) fn parse_date(str: &str) -> Result<DateTime<FixedOffset>, chrono::ParseError> {
+    chrono::NaiveDate::parse_from_str(str, "%Y-%m-%d")
+        .map(|date| DateTime::<Utc>::from_utc(date.and_hms_opt(0, 0, 0).unwrap(), Utc))
+        .map(DateTime::<FixedOffset>::from)
+}
+
+pub(crate) fn parse_timestamptz(str: &str) -> Result<DateTime<FixedOffset>, chrono::ParseError> {
+    DateTime::parse_from_rfc3339(str)
+}
+
+pub(crate) fn parse_timestamp(str: &str) -> Result<DateTime<FixedOffset>, chrono::ParseError> {
+    NaiveDateTime::parse_from_str(str, "%Y-%m-%dT%H:%M:%S%.f")
+        .map(|dt| DateTime::from_utc(dt, Utc))
+        .or_else(|_| DateTime::parse_from_rfc3339(str).map(DateTime::<Utc>::from))
+        .map(DateTime::<FixedOffset>::from)
+}
+
+pub(crate) fn parse_time(str: &str) -> Result<DateTime<FixedOffset>, chrono::ParseError> {
+    chrono::NaiveTime::parse_from_str(str, "%H:%M:%S%.f")
+        .map(|time| {
+            let base_date = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
+
+            DateTime::<Utc>::from_utc(base_date.and_time(time), Utc)
+        })
+        .map(DateTime::<FixedOffset>::from)
+}
+
+pub(crate) fn parse_money(str: &str) -> Result<BigDecimal, ParseBigDecimalError> {
+    // We strip out the currency sign from the string.
+    BigDecimal::from_str(&str[1..]).map(|bd| bd.normalized())
+}
+
+pub(crate) fn parse_decimal(str: &str) -> Result<BigDecimal, ParseBigDecimalError> {
+    BigDecimal::from_str(str).map(|bd| bd.normalized())
+}
diff --git a/psl/psl-core/src/datamodel_connector.rs b/psl/psl-core/src/datamodel_connector.rs
@@ -25,6 +25,8 @@ pub use self::{
 };
 
 use crate::{configuration::DatasourceConnectorData, Configuration, Datasource, PreviewFeature};
+use bigdecimal::{BigDecimal, ParseBigDecimalError};
+use chrono::{DateTime, FixedOffset};
 use diagnostics::{DatamodelError, Diagnostics, NativeTypeErrorFactory, Span};
 use enumflags2::BitFlags;
 use lsp_types::CompletionList;
@@ -359,6 +361,22 @@ pub trait Connector: Send + Sync {
     ) -> DatasourceConnectorData {
         Default::default()
     }
+
+    fn coerce_json_datetime(
+        &self,
+        _str: &str,
+        _nt: Option<NativeTypeInstance>,
+    ) -> chrono::ParseResult<DateTime<FixedOffset>> {
+        unreachable!("This method is only implemented on connectors with lateral join support.")
+    }
+
+    fn coerce_json_decimal(
+        &self,
+        _str: &str,
+        _nt: Option<NativeTypeInstance>,
+    ) -> Result<BigDecimal, ParseBigDecimalError> {
+        unreachable!("This method is only implemented on connectors with lateral join support.")
+    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq)]

diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/data_types/mod.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/data_types/mod.rs
@@ -7,5 +7,6 @@ mod enum_type;
 mod float;
 mod int;
 mod json;
+mod native;
 mod string;
 mod through_relation;
diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/data_types/native/mod.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/data_types/native/mod.rs
@@ -0,0 +1 @@
+mod postgres;