diff --git a/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs b/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs index e224cc52b15..94f29b21535 100644 --- a/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs +++ b/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs @@ -96,7 +96,7 @@ impl WriteOperations for MongoDbConnection { model: &Model, record_filter: connector_interface::RecordFilter, args: WriteArgs, - limit: Option, + limit: Option, _traceparent: Option, ) -> connector_interface::Result { catch(async move { @@ -121,7 +121,7 @@ impl WriteOperations for MongoDbConnection { _record_filter: connector_interface::RecordFilter, _args: WriteArgs, _selected_fields: FieldSelection, - _limit: Option, + _limit: Option, _traceparent: Option, ) -> connector_interface::Result { unimplemented!() @@ -164,7 +164,7 @@ impl WriteOperations for MongoDbConnection { &mut self, model: &Model, record_filter: connector_interface::RecordFilter, - limit: Option, + limit: Option, _traceparent: Option, ) -> connector_interface::Result { catch(write::delete_records( diff --git a/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs b/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs index 4b6c41364e1..31943e0dd6c 100644 --- a/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs +++ b/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs @@ -127,7 +127,7 @@ impl WriteOperations for MongoDbTransaction<'_> { model: &Model, record_filter: connector_interface::RecordFilter, args: connector_interface::WriteArgs, - limit: Option, + limit: Option, _traceparent: Option, ) -> connector_interface::Result { catch(async move { @@ -151,7 +151,7 @@ impl WriteOperations for MongoDbTransaction<'_> { _record_filter: connector_interface::RecordFilter, _args: connector_interface::WriteArgs, _selected_fields: FieldSelection, - _limit: Option, + _limit: Option, _traceparent: Option, ) -> connector_interface::Result { unimplemented!() @@ -193,7 +193,7 @@ impl WriteOperations for MongoDbTransaction<'_> { &mut self, model: &Model, record_filter: connector_interface::RecordFilter, - limit: Option, + limit: Option, _traceparent: Option, ) -> connector_interface::Result { catch(write::delete_records( diff --git a/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs b/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs index 537fbb1218c..3445d7f36d3 100644 --- a/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs +++ b/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs @@ -1,4 +1,5 @@ use super::*; +use crate::error::MongoError::ConversionError; use crate::{ error::{DecorateErrorWithFieldInformationExtension, MongoError}, filter::{FilterPrefix, MongoFilter, MongoFilterVisitor}, @@ -161,7 +162,7 @@ pub async fn update_records<'conn>( selectors .into_iter() .take(match update_type { - UpdateType::Many { limit } => limit.unwrap_or(i64::MAX), + UpdateType::Many { limit } => limit.unwrap_or(usize::MAX), UpdateType::One => 1, } as usize) .map(|p| { @@ -232,7 +233,7 @@ pub async fn delete_records<'conn>( session: &mut ClientSession, model: &Model, record_filter: RecordFilter, - limit: Option, + limit: Option, ) -> crate::Result { let coll = database.collection::(model.db_name()); let id_field = pick_singular_id(model); @@ -240,7 +241,7 @@ pub async fn delete_records<'conn>( let ids = if let Some(selectors) = record_filter.selectors { selectors .into_iter() - .take(limit.unwrap_or(i64::MAX) as usize) + .take(limit.unwrap_or(usize::MAX)) .map(|p| { (&id_field, p.values().next().unwrap()) .into_bson() @@ -309,7 +310,7 @@ async fn find_ids( session: &mut ClientSession, model: &Model, filter: MongoFilter, - limit: Option, + limit: Option, ) -> crate::Result> { let id_field = model.primary_identifier(); let mut builder = MongoReadQueryBuilder::new(model.clone()); @@ -325,7 +326,17 @@ async fn find_ids( let mut builder = builder.with_model_projection(id_field)?; - builder.limit = limit; + if let Some(limit) = limit { + builder.limit = match i64::try_from(limit) { + Ok(limit) => Some(limit), + Err(_) => { + return Err(ConversionError { + from: "usize".to_owned(), + to: "i64".to_owned(), + }) + } + } + } let query = builder.build()?; let docs = query.execute(collection, session).await?; diff --git a/query-engine/connectors/query-connector/src/interface.rs b/query-engine/connectors/query-connector/src/interface.rs index a3c958c99fb..3bf1614e039 100644 --- a/query-engine/connectors/query-connector/src/interface.rs +++ b/query-engine/connectors/query-connector/src/interface.rs @@ -286,7 +286,7 @@ pub trait WriteOperations { model: &Model, record_filter: RecordFilter, args: WriteArgs, - limit: Option, + limit: Option, traceparent: Option, ) -> crate::Result; @@ -300,7 +300,7 @@ pub trait WriteOperations { record_filter: RecordFilter, args: WriteArgs, selected_fields: FieldSelection, - limit: Option, + limit: Option, traceparent: Option, ) -> crate::Result; @@ -328,7 +328,7 @@ pub trait WriteOperations { &mut self, model: &Model, record_filter: RecordFilter, - limit: Option, + limit: Option, traceparent: Option, ) -> crate::Result; diff --git a/query-engine/connectors/query-connector/src/lib.rs b/query-engine/connectors/query-connector/src/lib.rs index e6631a7dcc2..c497f121ae9 100644 --- a/query-engine/connectors/query-connector/src/lib.rs +++ b/query-engine/connectors/query-connector/src/lib.rs @@ -20,6 +20,6 @@ pub type Result = std::result::Result; /// However when we updating any records we want to return an empty array if zero items were updated #[derive(PartialEq)] pub enum UpdateType { - Many { limit: Option }, + Many { limit: Option }, One, } diff --git a/query-engine/connectors/sql-query-connector/src/database/connection.rs b/query-engine/connectors/sql-query-connector/src/database/connection.rs index ae9ddfc8947..614f174e562 100644 --- a/query-engine/connectors/sql-query-connector/src/database/connection.rs +++ b/query-engine/connectors/sql-query-connector/src/database/connection.rs @@ -226,7 +226,7 @@ where model: &Model, record_filter: RecordFilter, args: WriteArgs, - limit: Option, + limit: Option, traceparent: Option, ) -> connector::Result { let ctx = Context::new(&self.connection_info, traceparent); @@ -243,7 +243,7 @@ where record_filter: RecordFilter, args: WriteArgs, selected_fields: FieldSelection, - limit: Option, + limit: Option, traceparent: Option, ) -> connector::Result { let ctx = Context::new(&self.connection_info, traceparent); @@ -274,7 +274,7 @@ where &mut self, model: &Model, record_filter: RecordFilter, - limit: Option, + limit: Option, traceparent: Option, ) -> connector::Result { let ctx = Context::new(&self.connection_info, traceparent); diff --git a/query-engine/connectors/sql-query-connector/src/database/operations/update.rs b/query-engine/connectors/sql-query-connector/src/database/operations/update.rs index d6817a72073..9ea13127a4f 100644 --- a/query-engine/connectors/sql-query-connector/src/database/operations/update.rs +++ b/query-engine/connectors/sql-query-connector/src/database/operations/update.rs @@ -104,7 +104,7 @@ pub(super) async fn update_many_from_filter( record_filter: RecordFilter, args: WriteArgs, selected_fields: Option<&ModelProjection>, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> crate::Result> { let update = build_update_and_set_query(model, args, None, ctx); @@ -133,7 +133,7 @@ pub(super) async fn update_many_from_ids_and_filter( record_filter: RecordFilter, args: WriteArgs, selected_fields: Option<&ModelProjection>, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> crate::Result<(Vec>, Vec)> { let filter_condition = FilterBuilder::without_top_level_joins().visit_filter(record_filter.filter.clone(), ctx); @@ -145,7 +145,7 @@ pub(super) async fn update_many_from_ids_and_filter( let updates = { let update = build_update_and_set_query(model, args, selected_fields, ctx); - let ids: Vec<&SelectionResult> = ids.iter().take(limit.unwrap_or(i64::MAX) as usize).collect(); + let ids: Vec<&SelectionResult> = ids.iter().take(limit.unwrap_or(usize::MAX)).collect(); chunk_update_with_ids(update, model, &ids, filter_condition, ctx)? }; diff --git a/query-engine/connectors/sql-query-connector/src/database/operations/write.rs b/query-engine/connectors/sql-query-connector/src/database/operations/write.rs index 2e2237e64a6..4ddee294b16 100644 --- a/query-engine/connectors/sql-query-connector/src/database/operations/write.rs +++ b/query-engine/connectors/sql-query-connector/src/database/operations/write.rs @@ -376,7 +376,7 @@ async fn generate_updates( record_filter: RecordFilter, args: WriteArgs, selected_fields: Option<&ModelProjection>, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> crate::Result>> { if record_filter.has_selectors() { @@ -399,7 +399,7 @@ pub(crate) async fn update_records( model: &Model, record_filter: RecordFilter, args: WriteArgs, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> crate::Result { if args.args.is_empty() { @@ -421,7 +421,7 @@ pub(crate) async fn update_records_returning( record_filter: RecordFilter, args: WriteArgs, selected_fields: FieldSelection, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> crate::Result { let field_names: Vec = selected_fields.db_names().collect(); @@ -458,7 +458,7 @@ pub(crate) async fn delete_records( conn: &dyn Queryable, model: &Model, record_filter: RecordFilter, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> crate::Result { let filter_condition = FilterBuilder::without_top_level_joins().visit_filter(record_filter.clone().filter, ctx); @@ -474,7 +474,8 @@ pub(crate) async fn delete_records( { row_count += conn.execute(delete).await?; if let Some(old_remaining_limit) = remaining_limit { - let new_remaining_limit = old_remaining_limit - row_count as i64; + // u64 to usize cast here cannot 'overflow' as the number of rows was limited to MAX usize in the first place. + let new_remaining_limit = old_remaining_limit - row_count as usize; if new_remaining_limit <= 0 { break; } diff --git a/query-engine/connectors/sql-query-connector/src/database/transaction.rs b/query-engine/connectors/sql-query-connector/src/database/transaction.rs index fde803502de..6528343f54f 100644 --- a/query-engine/connectors/sql-query-connector/src/database/transaction.rs +++ b/query-engine/connectors/sql-query-connector/src/database/transaction.rs @@ -220,7 +220,7 @@ impl WriteOperations for SqlConnectorTransaction<'_> { model: &Model, record_filter: RecordFilter, args: WriteArgs, - limit: Option, + limit: Option, traceparent: Option, ) -> connector::Result { let ctx = Context::new(&self.connection_info, traceparent); @@ -237,7 +237,7 @@ impl WriteOperations for SqlConnectorTransaction<'_> { record_filter: RecordFilter, args: WriteArgs, selected_fields: FieldSelection, - limit: Option, + limit: Option, traceparent: Option, ) -> connector::Result { let ctx = Context::new(&self.connection_info, traceparent); @@ -283,7 +283,7 @@ impl WriteOperations for SqlConnectorTransaction<'_> { &mut self, model: &Model, record_filter: RecordFilter, - limit: Option, + limit: Option, traceparent: Option, ) -> connector::Result { catch(&self.connection_info, async { diff --git a/query-engine/connectors/sql-query-connector/src/limit.rs b/query-engine/connectors/sql-query-connector/src/limit.rs index 281d538240a..8df1e749e4b 100644 --- a/query-engine/connectors/sql-query-connector/src/limit.rs +++ b/query-engine/connectors/sql-query-connector/src/limit.rs @@ -5,7 +5,7 @@ use query_structure::*; pub(crate) fn wrap_with_limit_subquery_if_needed<'a>( model: &Model, filter_condition: ConditionTree<'a>, - limit: Option, + limit: Option, ctx: &Context, ) -> ConditionTree<'a> { if let Some(limit) = limit { @@ -22,7 +22,7 @@ pub(crate) fn wrap_with_limit_subquery_if_needed<'a>( Select::from_table(model.as_table(ctx)) .columns(columns) .so_that(filter_condition) - .limit(limit as usize), + .limit(limit), ), ) } else { diff --git a/query-engine/connectors/sql-query-connector/src/query_builder/write.rs b/query-engine/connectors/sql-query-connector/src/query_builder/write.rs index bb6bd688237..2f4ab525e84 100644 --- a/query-engine/connectors/sql-query-connector/src/query_builder/write.rs +++ b/query-engine/connectors/sql-query-connector/src/query_builder/write.rs @@ -227,7 +227,7 @@ pub(crate) fn delete_returning( pub(crate) fn delete_many_from_filter( model: &Model, filter_condition: ConditionTree<'static>, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> Query<'static> { let filter_condition = wrap_with_limit_subquery_if_needed(model, filter_condition, limit, ctx); @@ -242,7 +242,7 @@ pub(crate) fn delete_many_from_ids_and_filter( model: &Model, ids: &[&SelectionResult], filter_condition: ConditionTree<'static>, - limit: Option, + limit: Option, ctx: &Context<'_>, ) -> Vec> { let columns: Vec<_> = ModelProjection::from(model.primary_identifier()) diff --git a/query-engine/core/src/query_ast/write.rs b/query-engine/core/src/query_ast/write.rs index 00063b0dcf7..ca0287179e3 100644 --- a/query-engine/core/src/query_ast/write.rs +++ b/query-engine/core/src/query_ast/write.rs @@ -368,7 +368,7 @@ pub struct UpdateManyRecords { /// Fields of updated records that client has requested to return. /// `None` if the connector does not support returning the updated rows. pub selected_fields: Option, - pub limit: Option, + pub limit: Option, } #[derive(Debug, Clone)] @@ -398,7 +398,7 @@ pub struct DeleteRecordFields { pub struct DeleteManyRecords { pub model: Model, pub record_filter: RecordFilter, - pub limit: Option, + pub limit: Option, } #[derive(Debug, Clone)] diff --git a/query-engine/core/src/query_graph_builder/write/delete.rs b/query-engine/core/src/query_graph_builder/write/delete.rs index b0c43d3d5ad..aa7d6e83205 100644 --- a/query-engine/core/src/query_graph_builder/write/delete.rs +++ b/query-engine/core/src/query_graph_builder/write/delete.rs @@ -1,12 +1,12 @@ use super::*; -use crate::query_document::ParsedInputValue; +use crate::query_graph_builder::write::limit::validate_limit; use crate::{ query_ast::*, query_graph::{Node, QueryGraph, QueryGraphDependency}, ArgumentListLookup, FilteredQuery, ParsedField, }; use psl::datamodel_connector::ConnectorCapability; -use query_structure::{Filter, Model, PrismaValue}; +use query_structure::{Filter, Model}; use schema::{constants::args, QuerySchema}; use std::convert::TryInto; @@ -111,13 +111,11 @@ pub fn delete_many_records( Some(where_arg) => extract_filter(where_arg.value.try_into()?, &model)?, None => Filter::empty(), }; - let limit = field - .arguments - .lookup(args::LIMIT) - .and_then(|limit_arg| match limit_arg.value { - ParsedInputValue::Single(PrismaValue::Int(i)) => Some(i), - _ => None, - }); + + let limit = match validate_limit(field.arguments.lookup(args::LIMIT)) { + Ok(limit) => limit, + Err(err) => return Err(err), + }; let model_id = model.primary_identifier(); let record_filter = filter.clone().into(); diff --git a/query-engine/core/src/query_graph_builder/write/limit.rs b/query-engine/core/src/query_graph_builder/write/limit.rs new file mode 100644 index 00000000000..2a615173ef9 --- /dev/null +++ b/query-engine/core/src/query_graph_builder/write/limit.rs @@ -0,0 +1,31 @@ +use crate::query_document::{ParsedArgument, ParsedInputValue}; +use crate::query_graph_builder::{QueryGraphBuilderError, QueryGraphBuilderResult}; +use query_structure::PrismaValue; + +pub(crate) fn validate_limit<'a>(limit_arg: Option>) -> QueryGraphBuilderResult> { + let limit = limit_arg.and_then(|limit_arg| match limit_arg.value { + ParsedInputValue::Single(PrismaValue::Int(i)) => Some(i), + _ => None, + }); + + match limit { + Some(i) => { + if i < 0 { + return Err(QueryGraphBuilderError::InputError(format!( + "Provided limit ({}) must be a positive integer.", + i + ))); + } + + match usize::try_from(i) { + Ok(i) => Ok(Some(i)), + Err(_) => Err(QueryGraphBuilderError::InputError(format!( + "Provided limit ({}) is beyond max int value for platform ({}).", + i, + usize::MAX + ))), + } + } + None => Ok(None), + } +} diff --git a/query-engine/core/src/query_graph_builder/write/mod.rs b/query-engine/core/src/query_graph_builder/write/mod.rs index 8db664e91d6..d1b1e62d0c5 100644 --- a/query-engine/core/src/query_graph_builder/write/mod.rs +++ b/query-engine/core/src/query_graph_builder/write/mod.rs @@ -2,6 +2,7 @@ mod connect; mod create; mod delete; mod disconnect; +mod limit; mod nested; mod raw; mod update; diff --git a/query-engine/core/src/query_graph_builder/write/update.rs b/query-engine/core/src/query_graph_builder/write/update.rs index a1e416f8079..db828599b76 100644 --- a/query-engine/core/src/query_graph_builder/write/update.rs +++ b/query-engine/core/src/query_graph_builder/write/update.rs @@ -1,5 +1,5 @@ use super::*; -use crate::query_document::ParsedInputValue; +use crate::query_graph_builder::write::limit::validate_limit; use crate::query_graph_builder::write::write_args_parser::*; use crate::ParsedObject; use crate::{ @@ -8,7 +8,7 @@ use crate::{ ArgumentListLookup, ParsedField, ParsedInputMap, }; use psl::datamodel_connector::ConnectorCapability; -use query_structure::{Filter, IntoFilter, Model, PrismaValue}; +use query_structure::{Filter, IntoFilter, Model}; use schema::{constants::args, QuerySchema}; use std::convert::TryInto; @@ -138,13 +138,10 @@ pub fn update_many_records( }; // "limit" - let limit = field - .arguments - .lookup(args::LIMIT) - .and_then(|limit_arg| match limit_arg.value { - ParsedInputValue::Single(PrismaValue::Int(i)) => Some(i), - _ => None, - }); + let limit = match validate_limit(field.arguments.lookup(args::LIMIT)) { + Ok(limit) => limit, + Err(err) => return Err(err), + }; // "data" let data_argument = field.arguments.lookup(args::DATA).unwrap(); @@ -362,5 +359,5 @@ fn can_use_atomic_update( pub struct UpdateManyRecordNodeOptionals<'a> { pub name: Option, pub nested_field_selection: Option>, - pub limit: Option, + pub limit: Option, }