Add mixin to handle annotated errors

I looked at the CockRoachDB implementation (compatible with Postgres)

Essentially, the simplest check is to ensure the ORDER BY columns exist in the DISTINCT ON columns.

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

@@ -3,12 +3,14 @@ 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 {
@@ -19,15 +21,24 @@ internal abstract class DistinctOnExpressionMixin(node: ASTNode) :
     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 orderByList = (parent.parent as SqlCompoundSelectStmtImpl).orderingTermList
-    // todo leftmost only
-    distinctOnColumns.forEachIndexed { idx, col ->
-      if (!col.textMatches(orderByList[idx])) {
+
+    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.any { distinctOnCol -> distinctOnCol.textMatches(orderByCol) }) {
         annotationHolder.createErrorAnnotation(
-          element = col,
-          message = "DISTINCT ON expression(s) must match the leftmost ORDER BY expression(s)",
+          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

@@ -4,10 +4,40 @@ CREATE TABLE person (
     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

@@ -1 +1,3 @@
-Test.s line 11:20 - DISTINCT ON expression(s) must match the leftmost ORDER BY expression(s)
+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