Fix 5345 PostgreSql SELECT DISTINCT ON (#5353)

* Update DistinctOnExpressionMixin.kt

Update DistinctOnExpressionMixin.kt - expose missing tables

Query the columns or any tables prefixing columns in the distinctOn clause

As there could be one or more columns with the same table prefix, remove any duplicates

* Add tests

fixture tests - multiple tables and columns in DISTINCT ON
integration test - valid sql is generated and helps prevent regressions

* Remove unnecessary nullability

As the code is was using mapNotNull
Remove unnecessary null type parameter and not-null assertion operator

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)
+    }
+  }
 }