diff --git a/psl/builtin-connectors/src/postgres_datamodel_connector.rs b/psl/builtin-connectors/src/postgres_datamodel_connector.rs index 8fac79165c58..a1bb73f8bd13 100644 --- a/psl/builtin-connectors/src/postgres_datamodel_connector.rs +++ b/psl/builtin-connectors/src/postgres_datamodel_connector.rs @@ -65,7 +65,8 @@ const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(Connector NativeUpsert | InsertReturning | UpdateReturning | - RowIn + RowIn | + DistinctOn }); pub struct PostgresDatamodelConnector; diff --git a/psl/psl-core/src/common/preview_features.rs b/psl/psl-core/src/common/preview_features.rs index 544bf5b99164..ca9da322933d 100644 --- a/psl/psl-core/src/common/preview_features.rs +++ b/psl/psl-core/src/common/preview_features.rs @@ -43,6 +43,7 @@ features!( ConnectOrCreate, CreateMany, DataProxy, + DistinctOn, Deno, Distinct, DriverAdapters, @@ -82,6 +83,7 @@ features!( pub const ALL_PREVIEW_FEATURES: FeatureMap = FeatureMap { active: enumflags2::make_bitflags!(PreviewFeature::{ Deno + | DistinctOn | DriverAdapters | FullTextIndex | FullTextSearch diff --git a/psl/psl-core/src/datamodel_connector/capabilities.rs b/psl/psl-core/src/datamodel_connector/capabilities.rs index 1b3f557e6285..3c0c1bbe47ed 100644 --- a/psl/psl-core/src/datamodel_connector/capabilities.rs +++ b/psl/psl-core/src/datamodel_connector/capabilities.rs @@ -103,7 +103,8 @@ capabilities!( NativeUpsert, InsertReturning, UpdateReturning, - RowIn, // Connector supports (a, b) IN (c, d) expression. + RowIn, // Connector supports (a, b) IN (c, d) expression. + DistinctOn // Connector supports DB-level distinct (e.g. postgres) ); /// Contains all capabilities that the connector is able to serve. diff --git a/psl/psl/tests/config/generators.rs b/psl/psl/tests/config/generators.rs index f10a9bda3eae..2f0708b13dc9 100644 --- a/psl/psl/tests/config/generators.rs +++ b/psl/psl/tests/config/generators.rs @@ -258,7 +258,7 @@ fn nice_error_for_unknown_generator_preview_feature() { .unwrap_err(); let expectation = expect![[r#" - error: The preview feature "foo" is not known. Expected one of: deno, driverAdapters, fullTextIndex, fullTextSearch, metrics, multiSchema, postgresqlExtensions, tracing, views + error: The preview feature "foo" is not known. Expected one of: distinctOn, deno, driverAdapters, fullTextIndex, fullTextSearch, metrics, multiSchema, postgresqlExtensions, tracing, views --> schema.prisma:3  |   2 |  provider = "prisma-client-js" diff --git a/quaint/src/ast.rs b/quaint/src/ast.rs index dc634423014a..b2b96fc16654 100644 --- a/quaint/src/ast.rs +++ b/quaint/src/ast.rs @@ -49,7 +49,7 @@ pub use ordering::{IntoOrderDefinition, Order, OrderDefinition, Orderable, Order pub use over::*; pub use query::{Query, SelectQuery}; pub use row::Row; -pub use select::Select; +pub use select::{DistinctType, Select}; pub use table::*; pub use union::Union; pub use update::*; diff --git a/quaint/src/ast/select.rs b/quaint/src/ast/select.rs index 96d50ba645c5..0c798d5227a0 100644 --- a/quaint/src/ast/select.rs +++ b/quaint/src/ast/select.rs @@ -4,7 +4,7 @@ use std::borrow::Cow; /// A builder for a `SELECT` statement. #[derive(Debug, PartialEq, Clone, Default)] pub struct Select<'a> { - pub(crate) distinct: bool, + pub(crate) distinct: Option>, pub(crate) tables: Vec>, pub(crate) columns: Vec>, pub(crate) conditions: Option>, @@ -18,6 +18,12 @@ pub struct Select<'a> { pub(crate) comment: Option>, } +#[derive(Debug, PartialEq, Clone)] +pub enum DistinctType<'a> { + Default, + OnClause(Vec>), +} + impl<'a> From> for Expression<'a> { fn from(sel: Select<'a>) -> Expression<'a> { Expression { @@ -236,7 +242,12 @@ impl<'a> Select<'a> { /// # } /// ``` pub fn distinct(mut self) -> Self { - self.distinct = true; + self.distinct = Some(DistinctType::Default); + self + } + + pub fn distinct_on(mut self, columns: Vec>) -> Self { + self.distinct = Some(DistinctType::OnClause(columns)); self } diff --git a/quaint/src/visitor.rs b/quaint/src/visitor.rs index c205b49dd279..18dee697722f 100644 --- a/quaint/src/visitor.rs +++ b/quaint/src/visitor.rs @@ -234,9 +234,15 @@ pub trait Visitor<'a> { self.write("SELECT ")?; - if select.distinct { - self.write("DISTINCT ")?; - } + if let Some(distinct) = select.distinct { + match distinct { + DistinctType::Default => self.write("DISTINCT ")?, + DistinctType::OnClause(columns) => { + self.write("DISTINCT ON ")?; + self.surround_with("(", ") ", |ref mut s| s.visit_columns(columns))?; + } + } + }; if !select.tables.is_empty() { if select.columns.is_empty() { diff --git a/quaint/src/visitor/postgres.rs b/quaint/src/visitor/postgres.rs index 648b3f0dc1ec..749b752709d4 100644 --- a/quaint/src/visitor/postgres.rs +++ b/quaint/src/visitor/postgres.rs @@ -785,6 +785,19 @@ mod tests { assert_eq!(expected_sql, sql); } + #[test] + fn test_distinct_on() { + let expected_sql = "SELECT DISTINCT ON (\"bar\", \"foo\") \"bar\" FROM \"test\""; + let query = Select::from_table("test").column(Column::new("bar")).distinct_on(vec![ + Expression::from(Column::from("bar")), + Expression::from(Column::from("foo")), + ]); + + let (sql, _) = Postgres::build(query).unwrap(); + + assert_eq!(expected_sql, sql); + } + #[test] fn test_distinct_with_subquery() { let expected_sql = "SELECT DISTINCT (SELECT $1 FROM \"test2\"), \"bar\" FROM \"test\""; diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/distinct.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/distinct.rs index 55ab265d22eb..d1002323f0e3 100644 --- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/distinct.rs +++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/distinct.rs @@ -1,31 +1,73 @@ use query_engine_tests::*; +/// `distinct on` queries involving `orderBy` +/// are currently forced through in-memory handling as otherwise, +/// they fail with the following error message. This does not affect _all_ +/// cases of `distinct on` in conjunction with `orderBy`. +/// +/// ```sql +/// SELECT DISTINCT ON expressions must match initial ORDER BY expressions +/// ``` +/// +/// `distinct on` queries _not_ involving `orderBy` return differently ordered +/// result sets, hence we need to duplicate certain tests to track snapshots +/// for both the in-db pg results, and the in-mem result sets. #[test_suite(schema(schemas::user_posts))] mod distinct { use indoc::indoc; - use query_engine_tests::assert_query; + use query_engine_tests::run_query; #[connector_test] async fn empty_database(runner: Runner) -> TestResult<()> { - assert_query!( - runner, - "query { findManyUser(distinct: [first_name, last_name]) { id, first_name, last_name } }", - r#"{"data":{"findManyUser":[]}}"# + insta::assert_snapshot!( + run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { id, first_name, last_name } + }") + ), + @r###"{"data":{"findManyUser":[]}}"### ); Ok(()) } /// Regression test for not selecting the fields the distinct is performed on: https://github.com/prisma/prisma/issues/5969 - #[connector_test] - async fn no_panic(runner: Runner) -> TestResult<()> { + #[connector_test(only(Postgres))] + async fn no_panic_pg(runner: Runner) -> TestResult<()> { test_user(&runner, r#"{ id: 1, first_name: "Joe", last_name: "Doe", email: "1" }"#).await?; test_user(&runner, r#"{ id: 2, first_name: "Doe", last_name: "Joe", email: "2" }"#).await?; - assert_query!( - runner, - "query { findManyUser(distinct: [first_name, last_name]) { id } }", - r#"{"data":{"findManyUser":[{"id":1},{"id":2}]}}"# + insta::assert_snapshot!( + run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { id } + }") + ), + @r###"{"data":{"findManyUser":[{"id":2},{"id":1}]}}"### + ); + + Ok(()) + } + + /// Regression test for not selecting the fields the distinct is performed on: https://github.com/prisma/prisma/issues/5969 + #[connector_test(exclude(Postgres))] + async fn no_panic_other(runner: Runner) -> TestResult<()> { + test_user(&runner, r#"{ id: 1, first_name: "Joe", last_name: "Doe", email: "1" }"#).await?; + test_user(&runner, r#"{ id: 2, first_name: "Doe", last_name: "Joe", email: "2" }"#).await?; + + insta::assert_snapshot!( + run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { id } + }") + ), + @r###"{"data":{"findManyUser":[{"id":1},{"id":2}]}}"### ); Ok(()) @@ -36,17 +78,22 @@ mod distinct { test_user(&runner, r#"{ id: 1, first_name: "Joe", last_name: "Doe", email: "1" }"#).await?; test_user(&runner, r#"{ id: 2, first_name: "Joe", last_name: "Doe", email: "2" }"#).await?; - assert_query!( - runner, - "query { findManyUser(distinct: first_name) { id } }", - r#"{"data":{"findManyUser":[{"id":1}]}}"# + insta::assert_snapshot!( + run_query!( + &runner, + indoc!("{ + findManyUser(distinct: first_name) + { id } + }") + ), + @r###"{"data":{"findManyUser":[{"id":1}]}}"### ); Ok(()) } - #[connector_test] - async fn with_duplicates(runner: Runner) -> TestResult<()> { + #[connector_test(only(Postgres))] + async fn with_duplicates_pg(runner: Runner) -> TestResult<()> { test_user(&runner, r#"{ id: 1, first_name: "Joe", last_name: "Doe", email: "1" }"#).await?; test_user( &runner, @@ -55,17 +102,44 @@ mod distinct { .await?; test_user(&runner, r#"{ id: 3, first_name: "Joe", last_name: "Doe", email: "3" }"#).await?; - assert_query!( - runner, - "query { findManyUser(distinct: [first_name, last_name]) { id, first_name, last_name } }", - r#"{"data":{"findManyUser":[{"id":1,"first_name":"Joe","last_name":"Doe"},{"id":2,"first_name":"Hans","last_name":"Wurst"}]}}"# + insta::assert_snapshot!(run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { id, first_name, last_name } + }") + ), + @r###"{"data":{"findManyUser":[{"id":2,"first_name":"Hans","last_name":"Wurst"},{"id":1,"first_name":"Joe","last_name":"Doe"}]}}"### + ); + + Ok(()) + } + + #[connector_test(exclude(Postgres))] + async fn with_duplicates_other(runner: Runner) -> TestResult<()> { + test_user(&runner, r#"{ id: 1, first_name: "Joe", last_name: "Doe", email: "1" }"#).await?; + test_user( + &runner, + r#"{ id: 2, first_name: "Hans", last_name: "Wurst", email: "2" }"#, + ) + .await?; + test_user(&runner, r#"{ id: 3, first_name: "Joe", last_name: "Doe", email: "3" }"#).await?; + + insta::assert_snapshot!(run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { id, first_name, last_name } + }") + ), + @r###"{"data":{"findManyUser":[{"id":1,"first_name":"Joe","last_name":"Doe"},{"id":2,"first_name":"Hans","last_name":"Wurst"}]}}"### ); Ok(()) } #[connector_test] - async fn with_skip(runner: Runner) -> TestResult<()> { + async fn with_skip_basic(runner: Runner) -> TestResult<()> { test_user(&runner, r#"{ id: 1, first_name: "Joe", last_name: "Doe", email: "1" }"#).await?; test_user( &runner, @@ -74,10 +148,15 @@ mod distinct { .await?; test_user(&runner, r#"{ id: 3, first_name: "Joe", last_name: "Doe", email: "3" }"#).await?; - assert_query!( - runner, - "query { findManyUser(skip: 1, distinct: [first_name, last_name]) { id, first_name, last_name } }", - r#"{"data":{"findManyUser":[{"id":2,"first_name":"Hans","last_name":"Wurst"}]}}"# + insta::assert_snapshot!( + run_query!( + &runner, + indoc!("{ + findManyUser(skip: 1, distinct: [first_name, last_name]) + { id, first_name, last_name } + }") + ), + @r###"{"data":{"findManyUser":[{"id":2,"first_name":"Hans","last_name":"Wurst"}]}}"### ); Ok(()) @@ -93,10 +172,18 @@ mod distinct { .await?; test_user(&runner, r#"{ id: 3, first_name: "Joe", last_name: "Doe", email: "3" }"#).await?; - assert_query!( - runner, - "query { findManyUser(orderBy: { first_name: asc }, skip: 1, distinct: [first_name, last_name]) { first_name, last_name } }", - r#"{"data":{"findManyUser":[{"first_name":"Joe","last_name":"Doe"}]}}"# + insta::assert_snapshot!( + run_query!( + &runner, + indoc!("{ + findManyUser( + orderBy: { first_name: asc }, + skip: 1, + distinct: [first_name, last_name]) + { first_name, last_name } + }") + ), + @r###"{"data":{"findManyUser":[{"first_name":"Joe","last_name":"Doe"}]}}"### ); Ok(()) @@ -112,18 +199,50 @@ mod distinct { .await?; test_user(&runner, r#"{ id: 3, first_name: "Joe", last_name: "Doe", email: "3" }"#).await?; - assert_query!( - runner, - "query { findManyUser(orderBy: { id: desc }, distinct: [first_name, last_name]) { id, first_name, last_name } }", - r#"{"data":{"findManyUser":[{"id":3,"first_name":"Joe","last_name":"Doe"},{"id":2,"first_name":"Hans","last_name":"Wurst"}]}}"# + insta::assert_snapshot!(run_query!( + &runner, + indoc!("{ + findManyUser( + orderBy: { id: desc }, + distinct: [first_name, last_name]) + { id, first_name, last_name } + }") + ), + @r###"{"data":{"findManyUser":[{"id":3,"first_name":"Joe","last_name":"Doe"},{"id":2,"first_name":"Hans","last_name":"Wurst"}]}}"### ); Ok(()) } /// Mut return only distinct records for top record, and only for those the distinct relation records. - #[connector_test] - async fn nested_distinct(runner: Runner) -> TestResult<()> { + #[connector_test(only(Postgres))] + async fn nested_distinct_pg(runner: Runner) -> TestResult<()> { + nested_dataset(&runner).await?; + + // Returns Users 1, 3, 4, 5 top + // 1 => ["3", "1", "2"] + // 4 => ["1"] + // 3 => [] + // 5 => ["2", "3"] + insta::assert_snapshot!(run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { + id + posts(distinct: [title], orderBy: { id: asc }) { + title + } + }}") + ), + @r###"{"data":{"findManyUser":[{"id":1,"posts":[{"title":"3"},{"title":"1"},{"title":"2"}]},{"id":4,"posts":[{"title":"1"}]},{"id":3,"posts":[]},{"id":5,"posts":[{"title":"2"},{"title":"3"}]}]}}"### + ); + + Ok(()) + } + + #[connector_test(exclude(Postgres))] + async fn nested_distinct_other(runner: Runner) -> TestResult<()> { nested_dataset(&runner).await?; // Returns Users 1, 3, 4, 5 top @@ -131,17 +250,18 @@ mod distinct { // 3 => [] // 4 => ["1"] // 5 => ["2", "3"] - assert_query!( - runner, - indoc! {"{ - findManyUser(distinct: [first_name, last_name]) { - id - posts(distinct: [title], orderBy: { id: asc }) { - title - } - } - }"}, - r#"{"data":{"findManyUser":[{"id":1,"posts":[{"title":"3"},{"title":"1"},{"title":"2"}]},{"id":3,"posts":[]},{"id":4,"posts":[{"title":"1"}]},{"id":5,"posts":[{"title":"2"},{"title":"3"}]}]}}"# + insta::assert_snapshot!(run_query!( + &runner, + indoc!("{ + findManyUser(distinct: [first_name, last_name]) + { + id + posts(distinct: [title], orderBy: { id: asc }) { + title + } + }}") + ), + @r###"{"data":{"findManyUser":[{"id":1,"posts":[{"title":"3"},{"title":"1"},{"title":"2"}]},{"id":3,"posts":[]},{"id":4,"posts":[{"title":"1"}]},{"id":5,"posts":[{"title":"2"},{"title":"3"}]}]}}"### ); Ok(()) @@ -157,17 +277,20 @@ mod distinct { // 4 => ["1"] // 3 => [] // 2 => ["2", "1"] - assert_query!( - runner, - indoc! {"{ - findManyUser(distinct: [first_name, last_name], orderBy: { id: desc }) { - id - posts(distinct: [title], orderBy: { id: desc }) { - title - } - } - }"}, - r#"{"data":{"findManyUser":[{"id":5,"posts":[{"title":"2"},{"title":"3"}]},{"id":4,"posts":[{"title":"1"}]},{"id":3,"posts":[]},{"id":2,"posts":[{"title":"2"},{"title":"1"}]}]}}"# + insta::assert_snapshot!(run_query!( + &runner, + indoc! {"{ + findManyUser( + distinct: [first_name, last_name], + orderBy: { id: desc } + ) + { + id + posts(distinct: [title], orderBy: { id: desc }) { title } + } + }"} + ), + @r###"{"data":{"findManyUser":[{"id":5,"posts":[{"title":"2"},{"title":"3"}]},{"id":4,"posts":[{"title":"1"}]},{"id":3,"posts":[]},{"id":2,"posts":[{"title":"2"},{"title":"1"}]}]}}"### ); Ok(()) diff --git a/query-engine/connectors/sql-query-connector/src/query_builder/read.rs b/query-engine/connectors/sql-query-connector/src/query_builder/read.rs index 3aa91288ea90..b3fbb548b64c 100644 --- a/query-engine/connectors/sql-query-connector/src/query_builder/read.rs +++ b/query-engine/connectors/sql-query-connector/src/query_builder/read.rs @@ -106,6 +106,17 @@ impl SelectDefinition for QueryArguments { .iter() .fold(select_ast, |acc, o| acc.order_by(o.order_definition.clone())); + let select_ast = if let Some(distinct) = self.distinct { + let distinct_fields = ModelProjection::from(distinct) + .as_columns(ctx) + .map(Expression::from) + .collect_vec(); + + select_ast.distinct_on(distinct_fields) + } else { + select_ast + }; + match limit { Some(limit) => (select_ast.limit(limit as usize), aggregation_joins.columns), None => (select_ast, aggregation_joins.columns), diff --git a/query-engine/core/src/interpreter/query_interpreters/inmemory_record_processor.rs b/query-engine/core/src/interpreter/query_interpreters/inmemory_record_processor.rs index f4c0465412e2..110972e77f75 100644 --- a/query-engine/core/src/interpreter/query_interpreters/inmemory_record_processor.rs +++ b/query-engine/core/src/interpreter/query_interpreters/inmemory_record_processor.rs @@ -20,9 +20,14 @@ impl InMemoryRecordProcessor { /// Creates a new processor from the given query args. /// The original args will be modified to prevent db level processing. pub(crate) fn new_from_query_args(args: &mut QueryArguments) -> Self { - let processor = Self { args: args.clone() }; + let mut processor = Self { args: args.clone() }; + + if !args.requires_inmemory_distinct() { + processor.args.distinct = None; + } else { + args.distinct = None; + } - args.distinct = None; args.ignore_take = true; args.ignore_skip = true; @@ -50,7 +55,12 @@ impl InMemoryRecordProcessor { records }; - let records = self.apply_distinct(records); + let records = if self.requires_inmemory_distinct() { + self.apply_distinct(records) + } else { + records + }; + let mut records = self.apply_pagination(records); if self.needs_reversed_order() { diff --git a/query-engine/core/src/interpreter/query_interpreters/nested_read.rs b/query-engine/core/src/interpreter/query_interpreters/nested_read.rs index fa4dc7c6e529..8573e2853c45 100644 --- a/query-engine/core/src/interpreter/query_interpreters/nested_read.rs +++ b/query-engine/core/src/interpreter/query_interpreters/nested_read.rs @@ -11,6 +11,7 @@ pub(crate) async fn m2m( trace_id: Option, ) -> InterpretationResult<(ManyRecords, Option>)> { let processor = InMemoryRecordProcessor::new_from_query_args(&mut query.args); + let parent_field = &query.parent_field; let child_link_id = parent_field.related_field().linking_fields(); diff --git a/query-engine/core/src/interpreter/query_interpreters/read.rs b/query-engine/core/src/interpreter/query_interpreters/read.rs index 464ac6651677..cac53ef5900a 100644 --- a/query-engine/core/src/interpreter/query_interpreters/read.rs +++ b/query-engine/core/src/interpreter/query_interpreters/read.rs @@ -1,8 +1,7 @@ -use super::*; +use super::{inmemory_record_processor::InMemoryRecordProcessor, *}; use crate::{interpreter::InterpretationResult, query_ast::*, result_ast::*}; use connector::{self, error::ConnectorError, ConnectionLike, RelAggregationRow, RelAggregationSelection}; use futures::future::{BoxFuture, FutureExt}; -use inmemory_record_processor::InMemoryRecordProcessor; use query_structure::ManyRecords; use std::collections::HashMap; use user_facing_errors::KnownError; diff --git a/query-engine/query-structure/src/query_arguments.rs b/query-engine/query-structure/src/query_arguments.rs index f9c222d80dbe..16eaa189dd88 100644 --- a/query-engine/query-structure/src/query_arguments.rs +++ b/query-engine/query-structure/src/query_arguments.rs @@ -1,3 +1,5 @@ +use psl::{datamodel_connector::ConnectorCapability, PreviewFeature}; + use crate::*; /// `QueryArguments` define various constraints queried data should fulfill: @@ -70,7 +72,30 @@ impl QueryArguments { /// retrieved by the connector or if it requires the query engine to fetch a raw set /// of records and perform certain operations itself, in-memory. pub fn requires_inmemory_processing(&self) -> bool { - self.distinct.is_some() || self.contains_unstable_cursor() || self.contains_null_cursor() + self.contains_unstable_cursor() || self.contains_null_cursor() || self.requires_inmemory_distinct() + } + + pub fn requires_inmemory_distinct(&self) -> bool { + self.distinct.is_some() && !self.can_distinct_in_db() + } + + fn can_distinct_in_db(&self) -> bool { + let has_distinct_feature = self + .model() + .dm + .schema + .configuration + .preview_features() + .contains(PreviewFeature::DistinctOn); + + let connector_can_distinct_in_db = self + .model() + .dm + .schema + .connector + .has_capability(ConnectorCapability::DistinctOn); + + has_distinct_feature && connector_can_distinct_in_db && self.order_by.is_empty() } /// An unstable cursor is a cursor that is used in conjunction with an unstable (non-unique) combination of orderBys.