Fix 4580 Add PostgreSql Select Distinct On

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/PostgreSql.bnf

@@ -10,6 +10,7 @@
   parserImports=[
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.ADD"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.ABORT"
+    "static com.alecstrong.sql.psi.core.psi.SqlTypes.ALL"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.ALTER"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.ALWAYS"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.AS"
@@ -27,6 +28,7 @@
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.DEFAULT"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.DELETE"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.DESC"
+    "static com.alecstrong.sql.psi.core.psi.SqlTypes.DISTINCT"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.DO"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.DOT"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.DROP"
@@ -40,6 +42,7 @@
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.FROM"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.GENERATED"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.GROUP"
+    "static com.alecstrong.sql.psi.core.psi.SqlTypes.HAVING"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.ID"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.IF"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.IGNORE"
@@ -63,13 +66,15 @@
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.REPLACE"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.ROLLBACK"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.RP"
+    "static com.alecstrong.sql.psi.core.psi.SqlTypes.SELECT"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.SET"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.STRING"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.TO"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.TRUE"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.UNIQUE"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.UPDATE"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.USING"
+    "static com.alecstrong.sql.psi.core.psi.SqlTypes.VALUES"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.WHERE"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.WITH"
     "static com.alecstrong.sql.psi.core.psi.SqlTypes.WITHOUT"
@@ -91,6 +96,7 @@ overrides ::= type_name
   | extension_expr
   | extension_stmt
   | create_index_stmt
+  | select_stmt
 
 column_constraint ::= [ CONSTRAINT {identifier} ] (
   PRIMARY KEY [ ASC | DESC ] {conflict_clause} |
@@ -319,6 +325,18 @@ alter_table_alter_column ::= ALTER [COLUMN] {column_name}
   pin = 1
 }
 
+distinct_on_expr ::= DISTINCT ON LP {result_column} ( COMMA {result_column} ) * RP {
+  mixin = "app.cash.sqldelight.dialects.postgresql.grammar.mixins.DistinctOnExpressionMixin"
+  implements = "com.alecstrong.sql.psi.core.psi.SqlCompositeElement"
+}
+
+select_stmt ::= SELECT ( [ distinct_on_expr ] | [ DISTINCT | ALL ] ) {result_column} ( COMMA {result_column} ) * [ FROM {join_clause} ] [ WHERE <<expr '-1'>> ] [{group_by}] [HAVING <<expr '-1'>>] | VALUES {values_expression} ( COMMA {values_expression} ) * {
+  extends = "com.alecstrong.sql.psi.core.psi.impl.SqlSelectStmtImpl"
+  implements = "com.alecstrong.sql.psi.core.psi.SqlSelectStmt"
+  override = true
+  pin = 1
+}
+
 compound_select_stmt ::= [ {with_clause} ] {select_stmt}  ( {compound_operator} {select_stmt} ) * [ ORDER BY {ordering_term} ( COMMA {ordering_term} ) * ] [ LIMIT {limiting_term} ] [ ( OFFSET | COMMA ) {limiting_term} ] [ FOR UPDATE [ 'SKIP' 'LOCKED' ] ] {
   extends = "com.alecstrong.sql.psi.core.psi.impl.SqlCompoundSelectStmtImpl"
   implements = "com.alecstrong.sql.psi.core.psi.SqlCompoundSelectStmt"

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/mixins/DistinctOnExpressionMixin.kt

@@ -0,0 +1,46 @@
+package app.cash.sqldelight.dialects.postgresql.grammar.mixins
+
+import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlDistinctOnExpr
+import com.alecstrong.sql.psi.core.SqlAnnotationHolder
+import com.alecstrong.sql.psi.core.psi.QueryElement
+import com.alecstrong.sql.psi.core.psi.SqlColumnName
+import com.alecstrong.sql.psi.core.psi.SqlCompositeElementImpl
+import com.alecstrong.sql.psi.core.psi.SqlResultColumn
+import com.alecstrong.sql.psi.core.psi.SqlSelectStmt
+import com.alecstrong.sql.psi.core.psi.impl.SqlCompoundSelectStmtImpl
+import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+
+internal abstract class DistinctOnExpressionMixin(node: ASTNode) :
+  SqlCompositeElementImpl(node), PostgreSqlDistinctOnExpr {
+
+  private val distinctOnColumns get() = children.filterIsInstance<SqlResultColumn>()
+
+  override fun queryAvailable(child: PsiElement): Collection<QueryElement.QueryResult> {
+    return (parent as SqlSelectStmt).queryExposed()
+  }
+
+  // Some idea of the basic validation finds the ORDER BY columns in the DISTINCT ON
+  // https://github.com/cockroachdb/cockroach/blob/b994d025c678f495cb8b93044e35a8c59595bd78/pkg/sql/opt/optbuilder/distinct.go#L87
+  override fun annotate(annotationHolder: SqlAnnotationHolder) {
+    super.annotate(annotationHolder)
+
+    val orderByTerms = (parent.parent as SqlCompoundSelectStmtImpl).orderingTermList
+
+    val orderByColumnNames =
+      orderByTerms.mapNotNull { PsiTreeUtil.findChildOfType(it, SqlColumnName::class.java) }
+
+    val distinctOnColumnNames =
+      distinctOnColumns.mapNotNull { PsiTreeUtil.findChildOfType(it, SqlColumnName::class.java) }
+
+    orderByColumnNames.zip(distinctOnColumnNames) { orderByCol, _ ->
+      if (distinctOnColumnNames.none { distinctOnCol -> distinctOnCol.textMatches(orderByCol) }) {
+        annotationHolder.createErrorAnnotation(
+          element = orderByCol,
+          message = "SELECT DISTINCT ON expressions must match initial ORDER BY expressions",
+        )
+      }
+    }
+  }
+}

dialects/postgresql/src/testFixtures/resources/fixtures_postgresql/select-distinct-on/Test.s

@@ -0,0 +1,43 @@
+CREATE TABLE person (
+    id INTEGER PRIMARY KEY,
+    name TEXT,
+    created_at TIMESTAMPTZ
+);
+
+SELECT DISTINCT ON (name) *
+FROM person;
+
+SELECT DISTINCT ON (name) *
+FROM person
+ORDER BY name, created_at DESC;
+
+SELECT DISTINCT ON (id, name) id, name
+FROM person
+ORDER BY name DESC;
+
+SELECT DISTINCT ON (name, id) id, name, created_at
+FROM person
+ORDER BY id DESC;
+
+SELECT DISTINCT ON (name, id) id, name
+FROM person
+ORDER BY id, name ASC;
+
+SELECT DISTINCT ON (name, id) id, name
+FROM person
+ORDER BY id, name, created_at ASC;
+
+-- fail
+SELECT DISTINCT ON (name) *
+FROM person
+ORDER BY created_at DESC;
+
+-- fail
+SELECT DISTINCT ON (name, created_at) id, name, created_at
+FROM person
+ORDER BY id, name, created_at DESC;
+
+-- fail
+SELECT DISTINCT ON (name, id) id, name, created_at
+FROM person
+ORDER BY name, created_at, id DESC;

dialects/postgresql/src/testFixtures/resources/fixtures_postgresql/select-distinct-on/failure.txt

@@ -0,0 +1,3 @@
+Test.s line 33:9 - SELECT DISTINCT ON expressions must match initial ORDER BY expressions
+Test.s line 38:9 - SELECT DISTINCT ON expressions must match initial ORDER BY expressions
+Test.s line 43:15 - SELECT DISTINCT ON expressions must match initial ORDER BY expressions