Skip to content

Commit

Permalink
impr(quaint/qe): PostgreSQL DISTINCT ON (#4223)
Browse files Browse the repository at this point in the history
Added basic `distinct on` capability and preview feature for PostgreSQL
Affects: `read_many`, `m2m`, and `one2m`

Updated `queries/distinct` tests to track both current in-mem and postgres `distinct on` results
  • Loading branch information
Druue authored Dec 4, 2023
1 parent f9fa153 commit a0c366a
Show file tree
Hide file tree
Showing 14 changed files with 276 additions and 73 deletions.
3 changes: 2 additions & 1 deletion psl/builtin-connectors/src/postgres_datamodel_connector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,8 @@ const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(Connector
NativeUpsert |
InsertReturning |
UpdateReturning |
RowIn
RowIn |
DistinctOn
});

pub struct PostgresDatamodelConnector;
Expand Down
2 changes: 2 additions & 0 deletions psl/psl-core/src/common/preview_features.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ features!(
ConnectOrCreate,
CreateMany,
DataProxy,
DistinctOn,
Deno,
Distinct,
DriverAdapters,
Expand Down Expand Up @@ -82,6 +83,7 @@ features!(
pub const ALL_PREVIEW_FEATURES: FeatureMap = FeatureMap {
active: enumflags2::make_bitflags!(PreviewFeature::{
Deno
| DistinctOn
| DriverAdapters
| FullTextIndex
| FullTextSearch
Expand Down
3 changes: 2 additions & 1 deletion psl/psl-core/src/datamodel_connector/capabilities.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion psl/psl/tests/config/generators.rs
Original file line number Diff line number Diff line change
Expand Up @@ -258,7 +258,7 @@ fn nice_error_for_unknown_generator_preview_feature() {
.unwrap_err();

let expectation = expect![[r#"
[1;91merror[0m: [1mThe preview feature "foo" is not known. Expected one of: deno, driverAdapters, fullTextIndex, fullTextSearch, metrics, multiSchema, postgresqlExtensions, tracing, views[0m
[1;91merror[0m: [1mThe preview feature "foo" is not known. Expected one of: distinctOn, deno, driverAdapters, fullTextIndex, fullTextSearch, metrics, multiSchema, postgresqlExtensions, tracing, views[0m
--> schema.prisma:3
 | 
 2 |  provider = "prisma-client-js"
Expand Down
2 changes: 1 addition & 1 deletion quaint/src/ast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down
15 changes: 13 additions & 2 deletions quaint/src/ast/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>>,
Expand All @@ -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 {
Expand Down Expand Up @@ -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
}

Expand Down
12 changes: 9 additions & 3 deletions quaint/src/visitor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
13 changes: 13 additions & 0 deletions quaint/src/visitor/postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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\"";
Expand Down
Loading

0 comments on commit a0c366a

Please sign in to comment.