fix: batched queries with same args in diff order (#3398)

fixes prisma/prisma#16267

query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batch/select_one_compound.rs

@@ -192,6 +192,25 @@ mod compound_batch {
         Ok(())
     }
 
+    #[connector_test]
+    async fn two_equal_queries(runner: Runner) -> TestResult<()> {
+        create_test_data(&runner).await?;
+
+        let queries = vec![
+            r#"query {findUniqueArtist(where:{firstName_lastName:{firstName:"Musti",lastName:"Naukio"} }) {firstName lastName}}"#.to_string(),
+            r#"query {findUniqueArtist(where:{firstName_lastName:{lastName:"Naukio",firstName:"Musti"} }) {firstName lastName}}"#.to_string(),
+        ];
+
+        let batch_results = runner.batch(queries, false, None).await?;
+
+        insta::assert_snapshot!(
+            batch_results.to_string(),
+            @r###"{"batchResult":[{"data":{"findUniqueArtist":{"firstName":"Musti","lastName":"Naukio"}}},{"data":{"findUniqueArtist":{"firstName":"Musti","lastName":"Naukio"}}}]}"###
+        );
+
+        Ok(())
+    }
+
     fn should_batch_schema() -> String {
         let schema = indoc! {
             r#"model Artist {

query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batch/select_one_singular.rs

@@ -228,8 +228,20 @@ mod singular_batch {
         create_test_data(&runner).await?;
 
         let queries = vec![
-            r#"query { findUniqueArtist(where: { ArtistId: 1}) { Name }}"#.to_string(),
-            r#"query { findUniqueArtist(where: { ArtistId: 1}) { Name }}"#.to_string(),
+            r#"query { findUniqueArtist(where: { ArtistId: 1 }) { Name }}"#.to_string(),
+            r#"query { findUniqueArtist(where: { ArtistId: 1 }) { Name }}"#.to_string(),
+        ];
+
+        let batch_results = runner.batch(queries, false, None).await?;
+        insta::assert_snapshot!(
+            batch_results.to_string(),
+            @r###"{"batchResult":[{"data":{"findUniqueArtist":{"Name":"ArtistWithoutAlbums"}}},{"data":{"findUniqueArtist":{"Name":"ArtistWithoutAlbums"}}}]}"###
+        );
+
+        // With non unique filters
+        let queries = vec![
+            r#"query { findUniqueArtist(where: { ArtistId: 1, Name: "ArtistWithoutAlbums" }) { Name }}"#.to_string(),
+            r#"query { findUniqueArtist(where: { Name: "ArtistWithoutAlbums", ArtistId: 1 }) { Name }}"#.to_string(),
         ];
 
         let batch_results = runner.batch(queries, false, None).await?;

query-engine/core/src/query_document/mod.rs

@@ -35,6 +35,7 @@ use indexmap::IndexMap;
 use prisma_models::ModelRef;
 use schema::QuerySchemaRef;
 use schema_builder::constants::*;
+use std::collections::HashMap;
 
 pub type QueryParserResult<T> = std::result::Result<T, QueryParserError>;
 
@@ -156,7 +157,7 @@ impl BatchDocumentTransaction {
 
 #[derive(Debug, Clone)]
 pub struct CompactedDocument {
-    pub arguments: Vec<Vec<(String, QueryValue)>>,
+    pub arguments: Vec<HashMap<String, QueryValue>>,
     pub nested_selection: Vec<String>,
     pub operation: Operation,
     pub keys: Vec<String>,
@@ -242,14 +243,15 @@ impl CompactedDocument {
         // Saving the stub of the query name for later use.
         let name = selections[0].name().replacen("findUnique", "", 1);
 
-        // Convert the selections into a vector of arguments. This defines the
+        // Convert the selections into a map of arguments. This defines the
         // response order and how we fetch the right data from the response set.
-        let arguments: Vec<Vec<(String, QueryValue)>> = selections
+        let arguments: Vec<HashMap<String, QueryValue>> = selections
             .into_iter()
             .map(|mut sel| {
                 let where_obj = sel.pop_argument().unwrap().1.into_object().unwrap();
+                let filter_map: HashMap<String, QueryValue> = extract_filter(where_obj, model).into_iter().collect();
 
-                extract_filter(where_obj, model)
+                filter_map
             })
             .collect();
 

query-engine/core/src/response_ir/mod.rs

@@ -16,7 +16,7 @@ use crate::QueryValue;
 use indexmap::IndexMap;
 use prisma_models::PrismaValue;
 use serde::ser::{Serialize, SerializeMap, SerializeSeq, Serializer};
-use std::{fmt, sync::Arc};
+use std::{collections::HashMap, fmt, sync::Arc};
 
 pub use ir_serializer::*;
 pub use response::*;
@@ -44,16 +44,15 @@ impl List {
         self.len() == 0
     }
 
-    pub fn index_by(self, keys: &[String]) -> Vec<(Vec<QueryValue>, Map)> {
-        let mut map = Vec::with_capacity(self.len());
+    pub fn index_by(self, keys: &[String]) -> Vec<(HashMap<String, QueryValue>, Map)> {
+        let mut map: Vec<(HashMap<String, QueryValue>, Map)> = Vec::with_capacity(self.len());
 
         for item in self.into_iter() {
             let inner = item.into_map().unwrap();
-
-            let key: Vec<QueryValue> = keys
+            let key: HashMap<String, QueryValue> = keys
                 .iter()
-                .map(|key| inner.get(key).unwrap().clone().into_value().unwrap())
-                .map(QueryValue::from)
+                .map(|key| (key.clone(), inner.get(key).unwrap().clone().into_value().unwrap()))
+                .map(|(key, val)| (key, QueryValue::from(val)))
                 .collect();
 
             map.push((key, inner));

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

@@ -4,7 +4,7 @@ use futures::FutureExt;
 use indexmap::IndexMap;
 use query_core::{
     schema::QuerySchemaRef, BatchDocument, BatchDocumentTransaction, CompactedDocument, Item, Operation, QueryDocument,
-    QueryExecutor, QueryValue, ResponseData, TxId,
+    QueryExecutor, ResponseData, TxId,
 };
 use std::{fmt, panic::AssertUnwindSafe};
 
@@ -123,8 +123,29 @@ impl<'a> GraphQlHandler<'a> {
             Ok(Ok(response_data)) => {
                 let mut gql_response: GQLResponse = response_data.into();
 
-                // We find the response data and make a hash from the given unique keys.
-                let data = gql_response
+                // At this point, many findUnique queries were converted to a single findMany query and that query was run.
+                // This means we have a list of results and we need to map each result back to their original findUnique query.
+                // `data` is the piece of logic that allows us to do that mapping.
+                // It takes the findMany response and converts it to a map of arguments to result.
+                // Let's take an example. Given the following batched queries:
+                // [
+                //    findUnique(where: { id: 1, name: "Bob" }) { id name age },
+                //    findUnique(where: { id: 2, name: "Alice" }) { id name age }
+                // ]
+                // 1. This gets converted to: findMany(where: { OR: [{ id: 1, name: "Bob" }, { id: 2, name: "Alice" }] }) { id name age }
+                // 2. Say we get the following result back: [{ id: 1, name: "Bob", age: 18 }, { id: 2, name: "Alice", age: 27 }]
+                // 3. We know the inputted arguments are ["id", "name"]
+                // 4. So we go over the result and build the following list:
+                // [
+                //  ({ id: 1, name: "Bob" },   { id: 1, name: "Bob", age: 18 }),
+                //  ({ id: 2, name: "Alice" }, { id: 2, name: "Alice", age: 27 })
+                // ]
+                // 5. Now, given the original findUnique queries and that list, we can easily find back which arguments maps to which result
+                // [
+                //    findUnique(where: { id: 1, name: "Bob" }) { id name age } -> { id: 1, name: "Bob", age: 18 }
+                //    findUnique(where: { id: 2, name: "Alice" }) { id name age } -> { id: 2, name: "Alice", age: 27 }
+                // ]
+                let args_to_results = gql_response
                     .take_data(plural_name)
                     .unwrap()
                     .into_list()
@@ -134,13 +155,13 @@ impl<'a> GraphQlHandler<'a> {
                 let results: Vec<GQLResponse> = arguments
                     .into_iter()
                     .map(|args| {
-                        let vals: Vec<QueryValue> = args.into_iter().map(|(_, v)| v).collect();
                         let mut responses = GQLResponse::with_capacity(1);
 
+                        // This is step 5 of the comment above.
                         // Copying here is mandatory due to some of the queries
                         // might be repeated with the same arguments in the original
                         // batch. We need to give the same answer for both of them.
-                        match data.iter().find(|(k, _)| k == &vals) {
+                        match args_to_results.iter().find(|(a, _)| *a == args) {
                             Some((_, result)) => {
                                 // Filter out all the keys not selected in the
                                 // original query.