Skip to content

Commit

Permalink
fix(qe): querying full table on batched findUnique() (#4789)
Browse files Browse the repository at this point in the history 
* test 23343

* Fixed test batch_23343

Updated SelectionSet to allow for same field from compound index and extra field in findUnique Where.

Co-authored-by: Flavian Desverne <[email protected]>

* quick ref (#4795)

Co-authored-by: Flavian Desverne <[email protected]>

---------

Co-authored-by: Flavian Desverne <[email protected]>
Co-authored-by: Serhii Tatarintsev <[email protected]>
  • Loading branch information
@Druue @Weakky
3 people authored Mar 28, 2024
1 parent 60bda88 commit 446e407
Show file tree
Hide file tree
Showing 3 changed files with 181 additions and 88 deletions.
63 changes: 63 additions & 0 deletions ...ne/connector-test-kit-rs/query-engine-tests/tests/queries/batching/select_one_compound.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}"#)
Expand Down
32 changes: 17 additions & 15 deletions query-engine/core/src/query_document/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<T> = std::result::Result<T, ValidationError>;

#[derive(Debug)]
Expand Down Expand Up @@ -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.
Expand Down
174 changes: 101 additions & 73 deletions query-engine/core/src/query_document/selection.rs
View file
Original file line number Diff line number Diff line change
@@ -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);

Expand Down Expand Up @@ -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<Item = &str> + '_ {
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<ArgumentValue>),
Multi(Vec<Vec<Cow<'a, str>>>, Vec<Vec<ArgumentValue>>),
pub enum SelectionSet {
Single(QuerySingle),
Many(Vec<QueryFilters>),
Empty,
}

impl<'a> Default for SelectionSet<'a> {
fn default() -> Self {
Self::Empty
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct QuerySingle(String, Vec<ArgumentValue>);

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<Self> {
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<Cow<'a, str>>, 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<Self> {
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<QueryFilters>) -> 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<dyn Iterator<Item = &str> + '_> {
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<In<'a>> for ArgumentValue {
fn from(other: In<'a>) -> Self {
impl From<In> 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)
});
Expand All @@ -211,16 +238,17 @@ impl<'a> From<In<'a>> 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<String, ArgumentValue> = IndexMap::new();
argument.insert(key.clone().into_owned(), val);
let mut argument = IndexMap::new();

argument.insert(key.to_string(), val);
acc.or(argument)
});

Expand Down

0 comments on commit 446e407

Please sign in to comment.