Fix 5345 PostgreSql SELECT DISTINCT ON

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

@@ -2,11 +2,13 @@ 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.NamedElement
 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.SqlTableName
 import com.alecstrong.sql.psi.core.psi.impl.SqlCompoundSelectStmtImpl
 import com.intellij.lang.ASTNode
 import com.intellij.psi.PsiElement
@@ -19,7 +21,13 @@ internal abstract class DistinctOnExpressionMixin(node: ASTNode) :
   private val distinctOnColumns get() = children.filterIsInstance<SqlResultColumn>()
 
   override fun queryAvailable(child: PsiElement): Collection<QueryElement.QueryResult> {
-    return (parent as SqlSelectStmt).queryExposed()
+    val distinctOnColumnsWithTablePrefix: List<NamedElement> =
+      distinctOnColumns.mapNotNull { PsiTreeUtil.findChildOfType(it, SqlTableName::class.java) }
+    return if (distinctOnColumnsWithTablePrefix.isEmpty()) {
+      (parent as SqlSelectStmt).queryExposed()
+    } else {
+      distinctOnColumnsWithTablePrefix.flatMap { tableAvailable(child, it.name) }.associateBy { it.table }.values
+    }
   }
 
   // Some idea of the basic validation finds the ORDER BY columns in the DISTINCT ON

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

@@ -4,28 +4,38 @@ CREATE TABLE person (
     created_at TIMESTAMPTZ
 );
 
-SELECT DISTINCT ON (name) *
+SELECT DISTINCT ON (person.name) *
 FROM person;
 
 SELECT DISTINCT ON (name) *
-FROM person
-ORDER BY name, created_at DESC;
+FROM person;
 
-SELECT DISTINCT ON (id, name) id, name
-FROM person
-ORDER BY name DESC;
+CREATE TABLE student(
+    student_id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL
+);
 
-SELECT DISTINCT ON (name, id) id, name, created_at
-FROM person
-ORDER BY id DESC;
+CREATE TABLE grade(
+    grade_id INTEGER PRIMARY KEY,
+    student_id INTEGER REFERENCES student(student_id),
+    grade INT NOT NULL,
+    grade_date TIMESTAMP NOT NULL
+);
 
-SELECT DISTINCT ON (name, id) id, name
-FROM person
-ORDER BY id, name ASC;
+SELECT DISTINCT ON (grade.student_id) grade.*, student.*
+FROM grade
+JOIN student USING (student_id)
+ORDER BY grade.student_id, grade_date;
 
-SELECT DISTINCT ON (name, id) id, name
-FROM person
-ORDER BY id, name, created_at ASC;
+SELECT DISTINCT ON (grade.student_id, grade.grade_date) grade.*, student.*
+FROM grade
+JOIN student USING (student_id)
+ORDER BY grade.student_id, grade_date;
+
+SELECT DISTINCT ON (student_id) *
+FROM grade
+JOIN student USING (student_id)
+ORDER BY student_id, grade_date;
 
 -- fail
 SELECT DISTINCT ON (name) *

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

@@ -1,3 +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
+Test.s line 43:9 - SELECT DISTINCT ON expressions must match initial ORDER BY expressions
+Test.s line 48:9 - SELECT DISTINCT ON expressions must match initial ORDER BY expressions
+Test.s line 53:15 - SELECT DISTINCT ON expressions must match initial ORDER BY expressions

sqldelight-gradle-plugin/src/test/integration-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/DistinctOn.sq

@@ -0,0 +1,29 @@
+CREATE TABLE student(
+    student_id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL
+);
+
+CREATE TABLE grade(
+    grade_id INTEGER PRIMARY KEY,
+    student_id INTEGER REFERENCES student(student_id),
+    grade INT NOT NULL,
+    grade_date TIMESTAMP NOT NULL
+);
+
+insertStudent:
+INSERT INTO student VALUES ?;
+
+insertGrade:
+INSERT INTO grade VALUES ?;
+
+selectDistinctOnStudent:
+SELECT DISTINCT ON (student_id) *
+FROM grade
+JOIN student USING (student_id)
+ORDER BY student_id, grade_date;
+
+selectDistinctOnStudentGradeDate:
+SELECT DISTINCT ON (grade.student_id, grade.grade_date) grade.*, student.*
+FROM grade
+JOIN student USING (student_id)
+ORDER BY grade.student_id, grade_date;

sqldelight-gradle-plugin/src/test/integration-postgresql/src/test/kotlin/app/cash/sqldelight/postgresql/integration/PostgreSqlTest.kt

@@ -1037,4 +1037,28 @@ class PostgreSqlTest {
       assertThat(x2).isEqualTo(b)
     }
   }
+
+  @Test
+  fun testSelectDistinctOn() {
+    val studentExpected = Student(1000, "Test Student")
+    val gradeExpected = Grade(4000, studentExpected.student_id, 5, LocalDateTime.of(1980, 1, 1, 1, 0, 0))
+    database.distinctOnQueries.insertStudent(studentExpected)
+    database.distinctOnQueries.insertGrade(gradeExpected)
+
+    with(database.distinctOnQueries.selectDistinctOnStudent().executeAsOne()) {
+      assertThat(student_id).isEqualTo(studentExpected.student_id)
+      assertThat(name).isEqualTo(studentExpected.name)
+      assertThat(grade_id).isEqualTo(gradeExpected.grade_id)
+      assertThat(grade).isEqualTo(gradeExpected.grade)
+      assertThat(grade_date).isEqualTo(gradeExpected.grade_date)
+    }
+
+    with(database.distinctOnQueries.selectDistinctOnStudentGradeDate().executeAsOne()) {
+      assertThat(student_id).isEqualTo(studentExpected.student_id)
+      assertThat(name).isEqualTo(studentExpected.name)
+      assertThat(grade_id).isEqualTo(gradeExpected.grade_id)
+      assertThat(grade).isEqualTo(gradeExpected.grade)
+      assertThat(grade_date).isEqualTo(gradeExpected.grade_date)
+    }
+  }
 }