diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batching/select_one_compound.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batching/select_one_compound.rs index ed94e4487bfe..0d055f591c72 100644 --- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batching/select_one_compound.rs +++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/batching/select_one_compound.rs @@ -422,6 +422,69 @@ mod compound_batch { Ok(()) } + fn schema_23343() -> String { + let schema = indoc! { r#" + model Post { + id Int + tenantId String + userId Int + text String + + @@unique([tenantId, userId]) + } + "# }; + + schema.to_owned() + } + + #[connector_test(schema(schema_23343))] + async fn batch_23343(runner: Runner) -> TestResult<()> { + create_test_data_23343(&runner).await?; + + let queries = vec![ + r#"query { + findUniquePost(where: { tenantId_userId: { tenantId: "tenant1", userId: 1 }, tenantId: "tenant1" }) + { id, tenantId, userId, text }}"# + .to_string(), + r#"query { + findUniquePost(where: { tenantId_userId: { tenantId: "tenant2", userId: 3 }, tenantId: "tenant2" }) + { id, tenantId, userId, text }}"# + .to_string(), + ]; + + let batch_results = runner.batch(queries, false, None).await?; + insta::assert_snapshot!( + batch_results.to_string(), + @r###"{"batchResult":[{"data":{"findUniquePost":{"id":1,"tenantId":"tenant1","userId":1,"text":"Post 1!"}}},{"data":{"findUniquePost":{"id":3,"tenantId":"tenant2","userId":3,"text":"Post 3!"}}}]}"### + ); + + Ok(()) + } + + async fn create_test_data_23343(runner: &Runner) -> TestResult<()> { + runner + .query(r#"mutation { createOnePost(data: { id: 1, tenantId: "tenant1", userId: 1, text: "Post 1!" }) { id } }"#) + .await? + .assert_success(); + + runner + .query(r#"mutation { createOnePost(data: { id: 2, tenantId: "tenant1", userId: 2, text: "Post 2!" }) { id } }"#) + .await? + .assert_success(); + + runner + .query(r#"mutation { createOnePost(data: { id: 3, tenantId: "tenant2", userId: 3, text: "Post 3!" }) { id } }"#) + .await? + .assert_success(); + + runner + .query(r#"mutation { createOnePost(data: { id: 4, tenantId: "tenant2", userId: 4, text: "Post 4!" }) { id } }"#) + .await? + .assert_success(); + + Ok(()) + } + async fn create_test_data(runner: &Runner) -> TestResult<()> { runner .query(r#"mutation { createOneArtist(data: { firstName: "Musti" lastName: "Naukio", non_unique: 0 }) { firstName }}"#) diff --git a/query-engine/core/src/query_document/mod.rs b/query-engine/core/src/query_document/mod.rs index fa424bc44d6e..575e3074df2f 100644 --- a/query-engine/core/src/query_document/mod.rs +++ b/query-engine/core/src/query_document/mod.rs @@ -37,6 +37,8 @@ use schema::{constants::*, QuerySchema}; use std::collections::HashMap; use user_facing_errors::query_engine::validation::ValidationError; +use self::selection::QueryFilters; + pub(crate) type QueryParserResult = std::result::Result; #[derive(Debug)] @@ -213,21 +215,21 @@ impl CompactedDocument { // The query arguments are extracted here. Combine all query // arguments from the different queries into a one large argument. - let selection_set = selections.iter().fold(SelectionSet::new(), |mut acc, selection| { - // findUnique always has only one argument. We know it must be an object, otherwise this will panic. - let where_obj = selection.arguments()[0] - .1 - .clone() - .into_object() - .expect("Trying to compact a selection with non-object argument"); - let filters = extract_filter(where_obj, &model); - - for (field, filter) in filters { - acc = acc.push(field, filter); - } - - acc - }); + let query_filters = selections + .iter() + .map(|selection| { + // findUnique always has only one argument. We know it must be an object, otherwise this will panic. + let where_obj = selection.arguments()[0] + .1 + .clone() + .into_object() + .expect("Trying to compact a selection with non-object argument"); + let filters = extract_filter(where_obj, &model); + + QueryFilters::new(filters) + }) + .collect(); + let selection_set = SelectionSet::new(query_filters); // We must select all unique fields in the query so we can // match the right response back to the right request later on. diff --git a/query-engine/core/src/query_document/selection.rs b/query-engine/core/src/query_document/selection.rs index 18f8fde78436..5b950fc38d3c 100644 --- a/query-engine/core/src/query_document/selection.rs +++ b/query-engine/core/src/query_document/selection.rs @@ -1,8 +1,9 @@ +use std::iter; + use crate::{ArgumentValue, ArgumentValueObject}; use indexmap::IndexMap; use itertools::Itertools; use schema::constants::filters; -use std::borrow::Cow; pub type SelectionArgument = (String, ArgumentValue); @@ -102,106 +103,132 @@ impl Selection { } } +#[derive(Debug, Clone, PartialEq, Default)] +pub struct QueryFilters(Vec<(String, ArgumentValue)>); + +impl QueryFilters { + pub fn new(filters: Vec<(String, ArgumentValue)>) -> Self { + Self(filters) + } + + pub fn keys(&self) -> impl IntoIterator + '_ { + self.0.iter().map(|(key, _)| key.as_str()) + } + + pub fn has_many_keys(&self) -> bool { + self.0.len() > 1 + } + + pub fn get_single_key(&self) -> Option<&(String, ArgumentValue)> { + self.0.first() + } +} + #[derive(Debug, Clone, PartialEq)] -pub enum SelectionSet<'a> { - Single(Cow<'a, str>, Vec), - Multi(Vec>>, Vec>), +pub enum SelectionSet { + Single(QuerySingle), + Many(Vec), Empty, } -impl<'a> Default for SelectionSet<'a> { - fn default() -> Self { - Self::Empty - } -} +#[derive(Debug, Clone, PartialEq)] +pub struct QuerySingle(String, Vec); + +impl QuerySingle { + /// Attempt at building a single query filter from multiple query filters. + /// Returns `None` if one of the query filters have more than one key. + pub fn new(query_filters: &[QueryFilters]) -> Option { + if query_filters.is_empty() { + return None; + } -impl<'a> SelectionSet<'a> { - pub fn new() -> Self { - Self::default() - } + if query_filters.iter().any(|query_filters| query_filters.has_many_keys()) { + return None; + } - pub fn push(self, column: impl Into>, value: ArgumentValue) -> Self { - let column = column.into(); + let first = query_filters.first().unwrap(); + let (key, value) = first.get_single_key().unwrap(); - match self { - Self::Single(key, mut vals) if key == column => { - vals.push(value); - Self::Single(key, vals) - } - Self::Single(key, mut vals) => { - vals.push(value); - Self::Multi(vec![vec![key, column]], vec![vals]) - } - Self::Multi(mut keys, mut vals) => { - match (keys.last_mut(), vals.last_mut()) { - (Some(keys), Some(vals)) if !keys.contains(&column) => { - keys.push(column); - vals.push(value); - } - _ => { - keys.push(vec![column]); - vals.push(vec![value]); - } - } + let mut result = QuerySingle(key.clone(), vec![value.clone()]); - Self::Multi(keys, vals) + for filters in query_filters.iter().skip(1) { + if let Some(single) = QuerySingle::push(result, filters) { + result = single; + } else { + return None; } - Self::Empty => Self::Single(column, vec![value]), } + + Some(result) } - pub fn len(&self) -> usize { - match self { - Self::Single(_, _) => 1, - Self::Multi(v, _) => v.len(), - Self::Empty => 0, + fn push(mut previous: Self, next: &QueryFilters) -> Option { + if next.0.is_empty() { + Some(previous) + // We have already validated that all `QueryFilters` have a single key. + // So we can continue building it. + } else { + let (key, value) = next.0.first().unwrap(); + + // if key matches, push value + if key == &previous.0 { + previous.1.push(value.clone()); + + Some(previous) + } else { + // if key does not match, it's a many + None + } } } +} - pub fn is_single(&self) -> bool { - matches!(self, Self::Single(_, _)) +impl Default for SelectionSet { + fn default() -> Self { + Self::Empty } +} - pub fn is_multi(&self) -> bool { - matches!(self, Self::Multi(_, _)) - } +impl SelectionSet { + pub fn new(filters: Vec) -> Self { + let single = QuerySingle::new(&filters); - pub fn is_empty(&self) -> bool { - self.len() == 0 + match single { + Some(single) => SelectionSet::Single(single), + None if filters.is_empty() => SelectionSet::Empty, + None => SelectionSet::Many(filters), + } } - pub fn keys(&self) -> Vec<&str> { + pub fn keys(&self) -> Box + '_> { match self { - Self::Single(key, _) => vec![key.as_ref()], - Self::Multi(keys, _) => match keys.first() { - Some(keys) => keys.iter().map(|key| key.as_ref()).collect(), - None => Vec::new(), - }, - Self::Empty => Vec::new(), + Self::Single(single) => Box::new(iter::once(single.0.as_str())), + Self::Many(filters) => Box::new(filters.iter().flat_map(|f| f.keys()).unique()), + Self::Empty => Box::new(iter::empty()), } } } -pub struct In<'a> { - selection_set: SelectionSet<'a>, +#[derive(Debug)] +pub struct In { + selection_set: SelectionSet, } -impl<'a> In<'a> { - pub fn new(selection_set: SelectionSet<'a>) -> Self { +impl In { + pub fn new(selection_set: SelectionSet) -> Self { Self { selection_set } } } -impl<'a> From> for ArgumentValue { - fn from(other: In<'a>) -> Self { +impl From for ArgumentValue { + fn from(other: In) -> Self { match other.selection_set { - SelectionSet::Multi(key_sets, val_sets) => { - let key_vals = key_sets.into_iter().zip(val_sets); - - let conjuctive = key_vals.fold(Conjuctive::new(), |acc, (keys, vals)| { - let ands = keys.into_iter().zip(vals).fold(Conjuctive::new(), |acc, (key, val)| { - let mut argument = IndexMap::new(); - argument.insert(key.into_owned(), val); + SelectionSet::Many(buckets) => { + let conjuctive = buckets.into_iter().fold(Conjuctive::new(), |acc, bucket| { + // Needed because we flush the last bucket by pushing an empty one, which gets translated to a `Null` as the Conjunctive is empty. + let ands = bucket.0.into_iter().fold(Conjuctive::new(), |acc, (key, value)| { + let mut argument = IndexMap::with_capacity(1); + argument.insert(key.clone(), value); acc.and(argument) }); @@ -211,16 +238,17 @@ impl<'a> From> for ArgumentValue { ArgumentValue::from(conjuctive) } - SelectionSet::Single(key, vals) => { - let is_bool = vals.iter().any(|v| match v { + SelectionSet::Single(QuerySingle(key, vals)) => { + let is_bool = vals.clone().into_iter().any(|v| match v { ArgumentValue::Scalar(s) => matches!(s, query_structure::PrismaValue::Boolean(_)), _ => false, }); if is_bool { let conjunctive = vals.into_iter().fold(Conjuctive::new(), |acc, val| { - let mut argument: IndexMap = IndexMap::new(); - argument.insert(key.clone().into_owned(), val); + let mut argument = IndexMap::new(); + + argument.insert(key.to_string(), val); acc.or(argument) });