Implement findOneOrThrow, findUniqueOrThrow

query-engine/core/src/interpreter/query_interpreters/read.rs

@@ -1,5 +1,5 @@
 use super::*;
-use crate::{interpreter::InterpretationResult, query_ast::*, result_ast::*};
+use crate::{interpreter::InterpretationResult, query_ast::*, result_ast::*, QueryGraphBuilderError};
 use connector::{self, ConnectionLike, QueryArguments, RelAggregationRow, RelAggregationSelection};
 use futures::future::{BoxFuture, FutureExt};
 use inmemory_record_processor::InMemoryRecordProcessor;
@@ -62,6 +62,14 @@ fn read_one(
                 .into())
             }
 
+            None if query
+                .options
+                .iter()
+                .any(|option| matches!(option, QueryOption::ThrowOnEmpty)) =>
+            {
+                Err(QueryGraphBuilderError::RecordNotFound("Record not found".to_owned()).into())
+            }
+
             None => Ok(QueryResult::RecordSelection(Box::new(RecordSelection {
                 name: query.name,
                 fields: query.selection_order,
@@ -120,18 +128,26 @@ fn read_many(
             (scalars, aggregation_rows)
         };
 
-        let nested: Vec<QueryResult> = process_nested(tx, query.nested, Some(&scalars)).await?;
-
-        Ok(RecordSelection {
-            name: query.name,
-            fields: query.selection_order,
-            scalars,
-            nested,
-            query_arguments: query.args,
-            model: query.model,
-            aggregation_rows,
+        if scalars.records.is_empty()
+            && query
+                .options
+                .iter()
+                .any(|option| matches!(option, QueryOption::ThrowOnEmpty))
+        {
+            Err(QueryGraphBuilderError::RecordNotFound("Record not found".to_owned()).into())
+        } else {
+            let nested: Vec<QueryResult> = process_nested(tx, query.nested, Some(&scalars)).await?;
+            Ok(RecordSelection {
+                name: query.name,
+                fields: query.selection_order,
+                scalars,
+                nested,
+                query_arguments: query.args,
+                model: query.model,
+                aggregation_rows,
+            }
+            .into())
         }
-        .into())
     };
 
     fut.boxed()
@@ -165,7 +181,6 @@ fn read_related<'conn>(
             )
             .await?
         };
-
         let model = query.parent_field.related_model();
         let nested: Vec<QueryResult> = process_nested(tx, query.nested, Some(&scalars)).await?;
 

query-engine/core/src/query_ast/read.rs

@@ -86,6 +86,14 @@ impl Display for ReadQuery {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum QueryOption {
+    ThrowOnEmpty,
+}
+
+pub const NO_QUERY_OPTIONS: [QueryOption; 0] = [];
+pub const THROW_ON_EMPTY: [QueryOption; 1] = [QueryOption::ThrowOnEmpty];
+
 #[derive(Debug, Clone)]
 pub struct RecordQuery {
     pub name: String,
@@ -96,6 +104,7 @@ pub struct RecordQuery {
     pub nested: Vec<ReadQuery>,
     pub selection_order: Vec<String>,
     pub aggregation_selections: Vec<RelAggregationSelection>,
+    pub options: Vec<QueryOption>,
 }
 
 #[derive(Debug, Clone)]
@@ -108,6 +117,7 @@ pub struct ManyRecordsQuery {
     pub nested: Vec<ReadQuery>,
     pub selection_order: Vec<String>,
     pub aggregation_selections: Vec<RelAggregationSelection>,
+    pub options: Vec<QueryOption>,
 }
 
 #[derive(Debug, Clone)]

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

@@ -818,6 +818,7 @@ impl QueryGraph {
                 nested: vec![],
                 selection_order: vec![],
                 aggregation_selections: vec![],
+                options: vec![],
             });
 
             let query = Query::Read(read_query);

query-engine/core/src/query_graph_builder/builder.rs

@@ -63,7 +63,9 @@ impl QueryGraphBuilder {
 
         let mut graph = match (&query_info.tag, query_info.model.clone()) {
             (QueryTag::FindUnique, Some(m)) => read::find_unique(parsed_field, m).map(Into::into),
+            (QueryTag::FindUniqueOrThrow, Some(m)) => read::find_unique_or_throw(parsed_field, m).map(Into::into),
             (QueryTag::FindFirst, Some(m)) => read::find_first(parsed_field, m).map(Into::into),
+            (QueryTag::FindFirstOrThrow, Some(m)) => read::find_first_or_throw(parsed_field, m).map(Into::into),
             (QueryTag::FindMany, Some(m)) => read::find_many(parsed_field, m).map(Into::into),
             (QueryTag::Aggregate, Some(m)) => read::aggregate(parsed_field, m).map(Into::into),
             (QueryTag::GroupBy, Some(m)) => read::group_by(parsed_field, m).map(Into::into),

query-engine/core/src/query_graph_builder/read/first.rs

@@ -4,14 +4,22 @@ use super::*;
 use crate::ParsedField;
 
 pub fn find_first(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
-    let mut many_query = many::find_many(field, model)?;
+    let many_query = many::find_many(field, model)?;
+    try_limit_to_one(many_query)
+}
+
+pub fn find_first_or_throw(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
+    let many_query = many::find_many_or_throw(field, model)?;
+    try_limit_to_one(many_query)
+}
 
-    // Optimization: Add `take: 1` to the query to reduce fetched result set size if possible.
-    Ok(match many_query {
+#[inline]
+fn try_limit_to_one(mut query: ReadQuery) -> QueryGraphBuilderResult<ReadQuery> {
+    Ok(match query {
         ReadQuery::ManyRecordsQuery(ref mut m) if m.args.take.is_none() => {
             m.args.take = Some(1);
-            many_query
+            query
         }
-        _ => many_query,
+        _ => query,
     })
 }

query-engine/core/src/query_graph_builder/read/many.rs

@@ -1,8 +1,21 @@
 use super::*;
-use crate::{query_document::ParsedField, ManyRecordsQuery, ReadQuery};
+use crate::{query_document::ParsedField, ManyRecordsQuery, QueryOption, ReadQuery, NO_QUERY_OPTIONS, THROW_ON_EMPTY};
 use prisma_models::ModelRef;
 
 pub fn find_many(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
+    find_many_with_options(field, model, NO_QUERY_OPTIONS.to_vec())
+}
+
+pub fn find_many_or_throw(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
+    find_many_with_options(field, model, THROW_ON_EMPTY.to_vec())
+}
+
+#[inline]
+fn find_many_with_options(
+    field: ParsedField,
+    model: ModelRef,
+    options: Vec<QueryOption>,
+) -> QueryGraphBuilderResult<ReadQuery> {
     let args = extractors::extract_query_args(field.arguments, &model)?;
     let name = field.name;
     let alias = field.alias;
@@ -26,5 +39,6 @@ pub fn find_many(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult
         nested,
         selection_order,
         aggregation_selections,
+        options,
     }))
 }

query-engine/core/src/query_graph_builder/read/one.rs

@@ -1,11 +1,24 @@
 use super::*;
-use crate::{query_document::*, ReadQuery, RecordQuery};
+use crate::{query_document::*, QueryOption, ReadQuery, RecordQuery, NO_QUERY_OPTIONS, THROW_ON_EMPTY};
 use prisma_models::ModelRef;
 use schema_builder::constants::args;
 use std::convert::TryInto;
 
+pub fn find_unique(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
+    find_unique_with_options(field, model, NO_QUERY_OPTIONS.to_vec())
+}
+
+pub fn find_unique_or_throw(field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
+    find_unique_with_options(field, model, THROW_ON_EMPTY.to_vec())
+}
+
 /// Builds a read query from a parsed incoming read query field.
-pub fn find_unique(mut field: ParsedField, model: ModelRef) -> QueryGraphBuilderResult<ReadQuery> {
+#[inline]
+pub fn find_unique_with_options(
+    mut field: ParsedField,
+    model: ModelRef,
+    options: Vec<QueryOption>,
+) -> QueryGraphBuilderResult<ReadQuery> {
     let filter = match field.arguments.lookup(args::WHERE) {
         Some(where_arg) => {
             let arg: ParsedInputMap = where_arg.value.try_into()?;
@@ -34,5 +47,6 @@ pub fn find_unique(mut field: ParsedField, model: ModelRef) -> QueryGraphBuilder
         nested,
         selection_order,
         aggregation_selections,
+        options,
     }))
 }

query-engine/core/src/query_graph_builder/write/utils.rs

@@ -44,6 +44,7 @@ where
         nested: vec![],
         selection_order: vec![],
         aggregation_selections: vec![],
+        options: vec![],
     });
 
     Query::Read(read_query)

query-engine/schema-builder/src/output_types/query_type.rs

@@ -9,12 +9,14 @@ pub(crate) fn build(ctx: &mut BuilderContext) -> (OutputType, ObjectTypeStrongRe
         .flat_map(|model| {
             let mut vec = vec![
                 find_first_field(ctx, &model),
+                find_first_or_throw_field(ctx, &model),
                 all_items_field(ctx, &model),
                 plain_aggregation_field(ctx, &model),
             ];
 
             vec.push(group_by_aggregation_field(ctx, &model));
             append_opt(&mut vec, find_unique_field(ctx, &model));
+            append_opt(&mut vec, find_unique_or_throw_field(ctx, &model));
 
             if ctx.enable_raw_queries && ctx.has_capability(ConnectorCapability::MongoDbQueryRaw) {
                 vec.push(mongo_find_raw_field(&model));
@@ -50,6 +52,25 @@ fn find_unique_field(ctx: &mut BuilderContext, model: &ModelRef) -> Option<Outpu
     })
 }
 
+/// Builds a "single" query arity item field (e.g. "user", "post" ...) for given model
+/// that will throw a NotFoundError if the item is not found
+fn find_unique_or_throw_field(ctx: &mut BuilderContext, model: &ModelRef) -> Option<OutputField> {
+    arguments::where_unique_argument(ctx, model).map(|arg| {
+        let field_name = format!("findUniqueOrThrow{}", model.name);
+
+        field(
+            field_name,
+            vec![arg],
+            OutputType::object(objects::model::map_type(ctx, model)),
+            Some(QueryInfo {
+                model: Some(Arc::clone(model)),
+                tag: QueryTag::FindUniqueOrThrow,
+            }),
+        )
+        .nullable()
+    })
+}
+
 /// Builds a find first item field for given model.
 fn find_first_field(ctx: &mut BuilderContext, model: &ModelRef) -> OutputField {
     let args = arguments::relation_selection_arguments(ctx, model, true);
@@ -67,6 +88,24 @@ fn find_first_field(ctx: &mut BuilderContext, model: &ModelRef) -> OutputField {
     .nullable()
 }
 
+/// Builds a find first item field for given model that throws a NotFoundError in case the item does
+/// not exist
+fn find_first_or_throw_field(ctx: &mut BuilderContext, model: &ModelRef) -> OutputField {
+    let args = arguments::relation_selection_arguments(ctx, model, true);
+    let field_name = format!("findFirstOrThrow{}", model.name);
+
+    field(
+        field_name,
+        args,
+        OutputType::object(objects::model::map_type(ctx, model)),
+        Some(QueryInfo {
+            model: Some(Arc::clone(model)),
+            tag: QueryTag::FindFirstOrThrow,
+        }),
+    )
+    .nullable()
+}
+
 /// Builds a "multiple" query arity items field (e.g. "users", "posts", ...) for given model.
 fn all_items_field(ctx: &mut BuilderContext, model: &ModelRef) -> OutputField {
     let args = arguments::relation_selection_arguments(ctx, model, true);

query-engine/schema/src/query_schema.rs

@@ -142,7 +142,9 @@ pub struct QueryInfo {
 #[derive(Debug, Clone, PartialEq)]
 pub enum QueryTag {
     FindUnique,
+    FindUniqueOrThrow,
     FindFirst,
+    FindFirstOrThrow,
     FindMany,
     CreateOne,
     CreateMany,
@@ -161,7 +163,9 @@ impl fmt::Display for QueryTag {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let s = match self {
             Self::FindUnique => "findUnique",
+            Self::FindUniqueOrThrow => "findUniqueOrThrow",
             Self::FindFirst => "findFirst",
+            Self::FindFirstOrThrow => "findFirstOrThrow",
             Self::FindMany => "findMany",
             Self::CreateOne => "createOne",
             Self::CreateMany => "createMany",