fix: allow combinations of keys of different types to make batching l…

…ogic consistent with single queries

query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batching/mod.rs

@@ -1,3 +1,4 @@
 mod select_one_compound;
 mod select_one_singular;
+mod select_with_diff_key_types;
 mod transactional_batch;

query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batching/select_with_diff_key_types.rs

@@ -0,0 +1,180 @@
+use query_engine_tests::*;
+
+#[test_suite(schema(schema))]
+mod select_different_key_type {
+    use indoc::indoc;
+    use query_engine_tests::{Runner, TestResult};
+
+    fn schema() -> String {
+        let schema = indoc! {
+            r#"
+                model FloatEntity {
+                    Id Float @id
+                    Text String
+                }
+
+                model BigIntEntity {
+                    Id BigInt @id
+                    Text String
+                }
+
+                model DecimalEntity {
+                    Id Decimal @id
+                    Text String
+                }
+
+                model DateEntity {
+                    Id DateTime @id
+                    Text String
+                }
+            "#
+        };
+
+        schema.to_owned()
+    }
+
+    #[connector_test]
+    async fn batch_of_two_distinct(runner: Runner) -> TestResult<()> {
+        create_test_data(&runner).await?;
+
+        // These pairs of queries are run as batched and non-batched to verify that the
+        // batching logic returns the same results as the non-batched logic.
+        // The reason this is valuable is that the batching logic, unlike regular queries
+        // relies on a comparison operation implemented in our code, which is sensitive
+        // to differences in types of values.
+        let pairs = [
+            (
+                r#"query { findUniqueFloatEntity(where: { Id: 1 }){ Text }}"#,
+                r#"query { findUniqueFloatEntity(where: { Id: 2 }){ Text }}"#,
+            ),
+            (
+                r#"query { findUniqueBigIntEntity(where: { Id: 1 }){ Text }}"#,
+                r#"query { findUniqueBigIntEntity(where: { Id: 2 }){ Text }}"#,
+            ),
+            (
+                r#"query { findUniqueDecimalEntity(where: { Id: 1 }){ Text }}"#,
+                r#"query { findUniqueDecimalEntity(where: { Id: 2 }){ Text }}"#,
+            ),
+            (
+                r#"query { findUniqueDateEntity(where: { Id: "2020-01-01T00:00:00Z" }){ Text }}"#,
+                r#"query { findUniqueDateEntity(where: { Id: "2020-01-02T00:00:00Z" }){ Text }}"#,
+            ),
+        ];
+
+        for (query_a, query_b) in pairs {
+            let batch = runner
+                .batch(vec![query_a.to_owned(), query_b.to_owned()], false, None)
+                .await?
+                .into_data();
+            let (single_a, single_b) = futures::try_join!(runner.query(query_a), runner.query(query_b),)?;
+
+            assert_eq!(batch.len(), 2, "{batch:?}");
+            assert_eq!(
+                batch,
+                single_a
+                    .into_data()
+                    .into_iter()
+                    .chain(single_b.into_data())
+                    .collect::<Vec<_>>()
+            );
+        }
+
+        Ok(())
+    }
+
+    #[connector_test]
+    async fn batch_of_two_repeated(runner: Runner) -> TestResult<()> {
+        create_test_data(&runner).await?;
+
+        // Same as above, but with just one row repeated twice in each batch.
+        let pairs = [
+            (
+                r#"query { findUniqueFloatEntity(where: { Id: 1 }){ Text }}"#,
+                r#"query { findUniqueFloatEntity(where: { Id: 1 }){ Text }}"#,
+            ),
+            (
+                r#"query { findUniqueBigIntEntity(where: { Id: 1 }){ Text }}"#,
+                r#"query { findUniqueBigIntEntity(where: { Id: 1 }){ Text }}"#,
+            ),
+            (
+                r#"query { findUniqueDecimalEntity(where: { Id: 1 }){ Text }}"#,
+                r#"query { findUniqueDecimalEntity(where: { Id: 1 }){ Text }}"#,
+            ),
+            (
+                r#"query { findUniqueDateEntity(where: { Id: "2020-01-01T00:00:00Z" }){ Text }}"#,
+                r#"query { findUniqueDateEntity(where: { Id: "2020-01-01T00:00:00Z" }){ Text }}"#,
+            ),
+        ];
+
+        for (query_a, query_b) in pairs {
+            let batch = runner
+                .batch(vec![query_a.to_owned(), query_b.to_owned()], false, None)
+                .await?
+                .into_data();
+            let (single_a, single_b) = futures::try_join!(runner.query(query_a), runner.query(query_b),)?;
+
+            assert_eq!(batch.len(), 2, "{batch:?}");
+            assert_eq!(
+                batch,
+                single_a
+                    .into_data()
+                    .into_iter()
+                    .chain(single_b.into_data())
+                    .collect::<Vec<_>>()
+            );
+        }
+
+        Ok(())
+    }
+
+    async fn create_test_data(runner: &Runner) -> TestResult<()> {
+        let mutations = [
+            r#"mutation entity {
+                createOneFloatEntity(data: { Id: 1, Text: "A" }) {
+                  Text
+                }
+              }"#,
+            r#"mutation entity {
+                createOneFloatEntity(data: { Id: 2, Text: "B" }) {
+                  Text
+                }
+              }"#,
+            r#"mutation entity {
+                createOneBigIntEntity(data: { Id: 1, Text: "A" }) {
+                  Text
+                }
+              }"#,
+            r#"mutation entity {
+                createOneBigIntEntity(data: { Id: 2, Text: "B" }) {
+                  Text
+                }
+            }"#,
+            r#"mutation entity {
+                createOneDecimalEntity(data: { Id: 1, Text: "A" }) {
+                  Text
+                }
+            }"#,
+            r#"mutation entity {
+                createOneDecimalEntity(data: { Id: 2, Text: "B" }) {
+                  Text
+                }
+            }"#,
+            r#"mutation entity {
+                createOneDateEntity(data: { Id: "2020-01-01T00:00:00Z", Text: "A" }) {
+                  Text
+                }
+            }"#,
+            r#"mutation entity {
+                createOneDateEntity(data: { Id: "2020-01-02T00:00:00Z", Text: "B" }) {
+                  Text
+                }
+            }"#,
+        ];
+
+        for mutation in mutations {
+            runner.query(mutation).await?.assert_success();
+        }
+
+        Ok(())
+    }
+}

query-engine/connector-test-kit-rs/query-tests-setup/src/query_result.rs

@@ -138,6 +138,14 @@ impl QueryResult {
         serde_json::to_string_pretty(&self.response).unwrap()
     }
 
+    pub fn into_data(self) -> Vec<serde_json::Value> {
+        match self.response {
+            Response::Single(res) => vec![res.data],
+            Response::Multi(res) => res.batch_result.into_iter().map(|res| res.data).collect(),
+            Response::Error(_) => vec![],
+        }
+    }
+
     /// Transform a JSON protocol response to a GraphQL protocol response, by removing the type
     /// tags.
     pub(crate) fn detag(&mut self) {

query-engine/request-handlers/src/handler.rs

@@ -255,6 +255,8 @@ impl<'a> RequestHandler<'a> {
     /// Here are the cases covered:
     /// - DateTime/String: User-input: DateTime / Response: String
     /// - Int/BigInt: User-input: Int / Response: BigInt
+    /// - Int/Float: User-input: Int / Response: Float
+    /// - Int/Decimal: User-input: Int / Response: String
     /// - (JSON protocol only) Custom types (eg: { "$type": "BigInt", value: "1" }): User-input: Scalar / Response: Object
     /// - (JSON protocol only) String/Enum: User-input: String / Response: Enum
     ///
@@ -267,6 +269,14 @@ impl<'a> RequestHandler<'a> {
                     .map(|t1| &t1 == t2)
                     .unwrap_or_else(|_| t1 == stringify_datetime(t2).as_str())
             }
+            (
+                ArgumentValue::Scalar(PrismaValue::Int(i1) | PrismaValue::BigInt(i1)),
+                ArgumentValue::Scalar(PrismaValue::Float(i2)),
+            )
+            | (
+                ArgumentValue::Scalar(PrismaValue::Float(i2)),
+                ArgumentValue::Scalar(PrismaValue::Int(i1) | PrismaValue::BigInt(i1)),
+            ) => BigDecimal::from(*i1) == *i2,
             (ArgumentValue::Scalar(PrismaValue::Int(i1)), ArgumentValue::Scalar(PrismaValue::BigInt(i2)))
             | (ArgumentValue::Scalar(PrismaValue::BigInt(i2)), ArgumentValue::Scalar(PrismaValue::Int(i1))) => {
                 *i1 == *i2
@@ -279,6 +289,16 @@ impl<'a> RequestHandler<'a> {
                 Some(t1) => Self::compare_values(t1, t2),
                 None => left == right,
             },
+            (
+                ArgumentValue::Scalar(PrismaValue::Int(s1) | PrismaValue::BigInt(s1)),
+                ArgumentValue::Scalar(PrismaValue::String(s2)),
+            )
+            | (
+                ArgumentValue::Scalar(PrismaValue::String(s2)),
+                ArgumentValue::Scalar(PrismaValue::Int(s1) | PrismaValue::BigInt(s1)),
+            ) => BigDecimal::from_str(s2)
+                .map(|s2| s2 == BigDecimal::from(*s1))
+                .unwrap_or(false),
             (ArgumentValue::Scalar(PrismaValue::Float(s1)), ArgumentValue::Scalar(PrismaValue::String(s2)))
             | (ArgumentValue::Scalar(PrismaValue::String(s2)), ArgumentValue::Scalar(PrismaValue::Float(s1))) => {
                 BigDecimal::from_str(s2).map(|s2| s2 == *s1).unwrap_or(false)