feat: generic query building

Cargo.toml

@@ -38,6 +38,7 @@ members = [
 
 [workspace.dependencies]
 async-trait = { version = "0.1.77" }
+bigdecimal = "0.3"
 enumflags2 = { version = "0.7", features = ["serde"] }
 futures = "0.3"
 psl = { path = "./psl/psl" }

query-engine/connectors/sql-query-connector/Cargo.toml

@@ -62,6 +62,9 @@ path = "../query-connector"
 [dependencies.query-structure]
 path = "../../query-structure"
 
+[dependencies.query-builder]
+path = "../../query-builders/query-builder"
+
 [dependencies.sql-query-builder]
 path = "../../query-builders/sql-query-builder"
 

query-engine/connectors/sql-query-connector/src/database/operations/read.rs

@@ -8,8 +8,9 @@ use crate::{QueryExt, Queryable, SqlError};
 use connector_interface::*;
 use futures::stream::{FuturesUnordered, StreamExt};
 use quaint::ast::*;
+use query_builder::QueryArgumentsExt;
 use query_structure::*;
-use sql_query_builder::{column_metadata, read, AsColumns, AsTable, Context, QueryArgumentsExt, RelationFieldExt};
+use sql_query_builder::{column_metadata, read, AsColumns, AsTable, Context, RelationFieldExt};
 
 pub(crate) async fn get_single_record(
     conn: &dyn Queryable,

query-engine/connectors/sql-query-connector/src/database/operations/read/process.rs

@@ -1,8 +1,8 @@
 use std::borrow::Cow;
 
 use itertools::{Either, Itertools};
+use query_builder::QueryArgumentsExt;
 use query_structure::{QueryArguments, Record};
-use sql_query_builder::QueryArgumentsExt;
 
 macro_rules! processor_state {
     ($name:ident $(-> $transition:ident($bound:ident))?) => {

query-engine/connectors/sql-query-connector/src/database/operations/write.rs

@@ -2,8 +2,7 @@ use super::update::*;
 use crate::row::ToSqlRow;
 use crate::value::to_prisma_value;
 use crate::{error::SqlError, QueryExt, Queryable};
-use itertools::Itertools;
-use quaint::ast::{Insert, Query};
+use quaint::ast::Query;
 use quaint::prelude::ResultSet;
 use quaint::{
     error::ErrorKind,
@@ -12,32 +11,9 @@ use quaint::{
 use query_structure::*;
 use sql_query_builder::{column_metadata, write, Context, FilterBuilder, SelectionResultExt, SqlTraceComment};
 use std::borrow::Cow;
-use std::{
-    collections::{HashMap, HashSet},
-    ops::Deref,
-};
+use std::collections::HashMap;
 use user_facing_errors::query_engine::DatabaseConstraint;
 
-#[cfg(target_arch = "wasm32")]
-macro_rules! trace {
-    (target: $target:expr, $($arg:tt)+) => {{
-        // No-op in WebAssembly
-    }};
-    ($($arg:tt)+) => {{
-        // No-op in WebAssembly
-    }};
-}
-
-#[cfg(not(target_arch = "wasm32"))]
-macro_rules! trace {
-    (target: $target:expr, $($arg:tt)+) => {
-        tracing::log::trace!(target: $target, $($arg)+);
-    };
-    ($($arg:tt)+) => {
-        tracing::log::trace!($($arg)+);
-    };
-}
-
 async fn generate_id(
     conn: &dyn Queryable,
     id_field: &FieldSelection,
@@ -191,49 +167,6 @@ pub(crate) async fn create_record(
     }
 }
 
-/// Returns a set of fields that are used in the arguments for the create operation.
-fn collect_affected_fields(args: &[WriteArgs], model: &Model) -> HashSet<ScalarFieldRef> {
-    let mut fields = HashSet::new();
-    args.iter().for_each(|arg| fields.extend(arg.keys()));
-
-    fields
-        .into_iter()
-        .map(|dsfn| model.fields().scalar().find(|sf| sf.db_name() == dsfn.deref()).unwrap())
-        .collect()
-}
-
-/// Generates a list of insert statements to execute. If `selected_fields` is set, insert statements
-/// will return the specified columns of inserted rows.
-pub fn generate_insert_statements(
-    model: &Model,
-    args: Vec<WriteArgs>,
-    skip_duplicates: bool,
-    selected_fields: Option<&ModelProjection>,
-    ctx: &Context<'_>,
-) -> Vec<Insert<'static>> {
-    let affected_fields = collect_affected_fields(&args, model);
-
-    if affected_fields.is_empty() {
-        args.into_iter()
-            .map(|_| write::create_records_empty(model, skip_duplicates, selected_fields, ctx))
-            .collect()
-    } else {
-        let partitioned_batches = partition_into_batches(args, ctx);
-        trace!("Total of {} batches to be executed.", partitioned_batches.len());
-        trace!(
-            "Batch sizes: {:?}",
-            partitioned_batches.iter().map(|b| b.len()).collect_vec()
-        );
-
-        partitioned_batches
-            .into_iter()
-            .map(|batch| {
-                write::create_records_nonempty(model, batch, skip_duplicates, &affected_fields, selected_fields, ctx)
-            })
-            .collect()
-    }
-}
-
 /// Inserts records specified as a list of `WriteArgs`. Returns number of inserted records.
 pub(crate) async fn create_records_count(
     conn: &dyn Queryable,
@@ -242,7 +175,7 @@ pub(crate) async fn create_records_count(
     skip_duplicates: bool,
     ctx: &Context<'_>,
 ) -> crate::Result<usize> {
-    let inserts = generate_insert_statements(model, args, skip_duplicates, None, ctx);
+    let inserts = write::generate_insert_statements(model, args, skip_duplicates, None, ctx);
     let mut count = 0;
     for insert in inserts {
         count += conn.execute(insert.into()).await?;
@@ -265,7 +198,7 @@ pub(crate) async fn create_records_returning(
     let idents = selected_fields.type_identifiers_with_arities();
     let meta = column_metadata::create(&field_names, &idents);
     let mut records = ManyRecords::new(field_names.clone());
-    let inserts = generate_insert_statements(model, args, skip_duplicates, Some(&selected_fields.into()), ctx);
+    let inserts = write::generate_insert_statements(model, args, skip_duplicates, Some(&selected_fields.into()), ctx);
 
     for insert in inserts {
         let result_set = conn.query(insert.into()).await?;
@@ -281,74 +214,6 @@ pub(crate) async fn create_records_returning(
     Ok(records)
 }
 
-/// Partitions data into batches, respecting `max_bind_values` and `max_insert_rows` settings from
-/// the `Context`.
-fn partition_into_batches(args: Vec<WriteArgs>, ctx: &Context<'_>) -> Vec<Vec<WriteArgs>> {
-    let batches = if let Some(max_params) = ctx.max_bind_values() {
-        // We need to split inserts if they are above a parameter threshold, as well as split based on number of rows.
-        // -> Horizontal partitioning by row number, vertical by number of args.
-        args.into_iter()
-            .peekable()
-            .batching(|iter| {
-                let mut param_count: usize = 0;
-                let mut batch = vec![];
-
-                while param_count < max_params {
-                    // If the param count _including_ the next item doens't exceed the limit,
-                    // we continue filling up the current batch.
-                    let proceed = match iter.peek() {
-                        Some(next) => (param_count + next.len()) <= max_params,
-                        None => break,
-                    };
-
-                    if proceed {
-                        match iter.next() {
-                            Some(next) => {
-                                param_count += next.len();
-                                batch.push(next)
-                            }
-                            None => break,
-                        }
-                    } else {
-                        break;
-                    }
-                }
-
-                if batch.is_empty() {
-                    None
-                } else {
-                    Some(batch)
-                }
-            })
-            .collect_vec()
-    } else {
-        vec![args]
-    };
-
-    if let Some(max_rows) = ctx.max_insert_rows() {
-        let capacity = batches.len();
-        batches
-            .into_iter()
-            .fold(Vec::with_capacity(capacity), |mut batches, next_batch| {
-                if next_batch.len() > max_rows {
-                    batches.extend(
-                        next_batch
-                            .into_iter()
-                            .chunks(max_rows)
-                            .into_iter()
-                            .map(|chunk| chunk.into_iter().collect_vec()),
-                    );
-                } else {
-                    batches.push(next_batch);
-                }
-
-                batches
-            })
-    } else {
-        batches
-    }
-}
-
 /// Update one record in a database defined in `conn` and the records
 /// defined in `args`, resulting the identifiers that were modified in the
 /// operation.

query-engine/connectors/sql-query-connector/src/lib.rs

@@ -11,8 +11,6 @@ mod value;
 use self::{query_ext::QueryExt, row::*};
 use quaint::prelude::Queryable;
 
-pub use database::operations::write::generate_insert_statements;
-
 pub use database::FromSource;
 #[cfg(feature = "driver-adapters")]
 pub use database::Js;

query-engine/core/Cargo.toml

@@ -9,7 +9,7 @@ graphql-protocol = []
 
 [dependencies]
 async-trait.workspace = true
-bigdecimal = "0.3"
+bigdecimal.workspace = true
 chrono.workspace = true
 connection-string.workspace = true
 connector = { path = "../connectors/query-connector", package = "query-connector" }

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

@@ -4,7 +4,12 @@ pub mod translate;
 use std::sync::Arc;
 
 pub use expression::Expression;
+use quaint::{
+    prelude::{ConnectionInfo, ExternalConnectionInfo, SqlFamily},
+    visitor::Postgres,
+};
 use schema::QuerySchema;
+use sql_query_builder::{Context, SqlQueryBuilder};
 use thiserror::Error;
 pub use translate::{translate, TranslateError};
 
@@ -23,7 +28,14 @@ pub fn compile(query_schema: &Arc<QuerySchema>, query_doc: QueryDocument) -> cra
     let QueryDocument::Single(query) = query_doc else {
         return Err(CompileError::UnsupportedRequest.into());
     };
+    let connection_info = ConnectionInfo::External(ExternalConnectionInfo::new(
+        SqlFamily::Postgres,
+        "public".to_owned(),
+        None,
+    ));
 
+    let ctx = Context::new(&connection_info, None);
+    let builder = SqlQueryBuilder::<Postgres<'_>>::new(ctx);
     let (graph, _serializer) = QueryGraphBuilder::new(query_schema).build(query)?;
-    Ok(translate(graph).map_err(CompileError::from)?)
+    Ok(translate(graph, &builder).map_err(CompileError::from)?)
 }

query-engine/core/src/compiler/translate.rs

@@ -1,6 +1,7 @@
 mod query;
 
 use query::translate_query;
+use query_builder::QueryBuilder;
 use thiserror::Error;
 
 use crate::{EdgeRef, Node, NodeRef, Query, QueryGraph};
@@ -12,17 +13,17 @@ pub enum TranslateError {
     #[error("node {0} has no content")]
     NodeContentEmpty(String),
 
-    #[error("{0}")]
-    QuaintError(#[from] quaint::error::Error),
+    #[error("query builder error: {0}")]
+    QueryBuildFailure(#[source] Box<dyn std::error::Error + Send + Sync>),
 }
 
 pub type TranslateResult<T> = Result<T, TranslateError>;
 
-pub fn translate(mut graph: QueryGraph) -> TranslateResult<Expression> {
+pub fn translate(mut graph: QueryGraph, builder: &dyn QueryBuilder) -> TranslateResult<Expression> {
     graph
         .root_nodes()
         .into_iter()
-        .map(|node| NodeTranslator::new(&mut graph, node, &[]).translate())
+        .map(|node| NodeTranslator::new(&mut graph, node, &[], builder).translate())
         .collect::<TranslateResult<Vec<_>>>()
         .map(Expression::Seq)
 }
@@ -32,14 +33,21 @@ struct NodeTranslator<'a, 'b> {
     node: NodeRef,
     #[allow(dead_code)]
     parent_edges: &'b [EdgeRef],
+    query_builder: &'b dyn QueryBuilder,
 }
 
 impl<'a, 'b> NodeTranslator<'a, 'b> {
-    fn new(graph: &'a mut QueryGraph, node: NodeRef, parent_edges: &'b [EdgeRef]) -> Self {
+    fn new(
+        graph: &'a mut QueryGraph,
+        node: NodeRef,
+        parent_edges: &'b [EdgeRef],
+        query_builder: &'b dyn QueryBuilder,
+    ) -> Self {
         Self {
             graph,
             node,
             parent_edges,
+            query_builder,
         }
     }
 
@@ -64,7 +72,7 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
             .try_into()
             .expect("current node must be query");
 
-        translate_query(query)
+        translate_query(query, self.query_builder)
     }
 
     #[allow(dead_code)]
@@ -99,7 +107,7 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
             .into_iter()
             .map(|(_, node)| {
                 let edges = self.graph.incoming_edges(&node);
-                NodeTranslator::new(self.graph, node, &edges).translate()
+                NodeTranslator::new(self.graph, node, &edges, self.query_builder).translate()
             })
             .collect::<Result<Vec<_>, _>>()?;
 
@@ -121,7 +129,7 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
             .map(|(_, node)| {
                 let name = node.id();
                 let edges = self.graph.incoming_edges(&node);
-                let expr = NodeTranslator::new(self.graph, node, &edges).translate()?;
+                let expr = NodeTranslator::new(self.graph, node, &edges, self.query_builder).translate()?;
                 Ok(Binding { name, expr })
             })
             .collect::<TranslateResult<Vec<_>>>()?;