impr(quaint/qe): PostgreSQL DISTINCT ON

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;

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.

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::*;

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<DistinctType<'a>>,
     pub(crate) tables: Vec<Table<'a>>,
     pub(crate) columns: Vec<Expression<'a>>,
     pub(crate) conditions: Option<ConditionTree<'a>>,
@@ -18,6 +18,12 @@ pub struct Select<'a> {
     pub(crate) comment: Option<Cow<'a, str>>,
 }
 
+#[derive(Debug, PartialEq, Clone)]
+pub enum DistinctType<'a> {
+    Default,
+    OnClause(Vec<Expression<'a>>),
+}
+
 impl<'a> From<Select<'a>> 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<Expression<'a>>) -> Self {
+        self.distinct = Some(DistinctType::OnClause(columns));
         self
     }
 

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() {

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\"";

query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/distinct.rs

@@ -3,14 +3,19 @@ use query_engine_tests::*;
 #[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(())
@@ -22,10 +27,15 @@ 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: "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(())
@@ -36,10 +46,15 @@ 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(())
@@ -55,17 +70,21 @@ 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]
-    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 +93,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 +117,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,10 +144,16 @@ 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(())
@@ -131,17 +169,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":4,"posts":[{"title":"1"}]},{"id":3,"posts":[]},{"id":5,"posts":[{"title":"2"},{"title":"3"}]}]}}"###
         );
 
         Ok(())
@@ -157,17 +196,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(())

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),