Fixes 5032 column adjacency for PostgreSql UPDATE FROM statement (#5035)

* Fixes column adjacency for Update statement Same fix for Sqlite * Add Integration tests * fix from spotless whitespace --------- Co-authored-by: Alec Kazakova <[email protected]>
cashapp · Apr 5, 2024 · b5d9d3f · b5d9d3f
1 parent aa41c0b
commit b5d9d3f
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 4 deletions.
diff --git a/...n/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/mixins/UpdateStmtLimitedMixin.kt b/...n/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/mixins/UpdateStmtLimitedMixin.kt
@@ -18,9 +18,9 @@ internal abstract class UpdateStmtLimitedMixin(
   FromQuery {
   override fun queryAvailable(child: PsiElement): Collection<QueryElement.QueryResult> {
     if (child != joinClause && joinClause != null) {
-      return super.queryAvailable(child) + joinClause!!.queryExposed()
+      return super.queryAvailable(child) +
+        joinClause!!.queryExposed().map { it.copy(adjacent = true) }
     }
-
     return super.queryAvailable(child)
   }
 

diff --git a/dialects/postgresql/src/testFixtures/resources/fixtures_postgresql/update-set-from/1.s b/dialects/postgresql/src/testFixtures/resources/fixtures_postgresql/update-set-from/1.s
@@ -1,9 +1,11 @@
 CREATE TABLE test(
-  id SERIAL PRIMARY KEY
+  id SERIAL PRIMARY KEY,
+  id2 INTEGER
 );
 
 CREATE TABLE test2(
-  id2 SERIAL PRIMARY KEY
+  id2 SERIAL PRIMARY KEY,
+  other TEXT
 );
 
 UPDATE test
@@ -40,4 +42,8 @@ FROM (
   ON otherTest.id = test2.id2
 );
 
+UPDATE test
+SET id2 = t2.id2
+FROM test2 t2
+WHERE other = 'x';
 
diff --git a/...tion-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/Updates.sq b/...tion-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/Updates.sq
@@ -0,0 +1,30 @@
+CREATE TABLE test(
+  id SERIAL PRIMARY KEY,
+  id2 INTEGER
+);
+
+CREATE TABLE test2(
+  id2 SERIAL PRIMARY KEY,
+  other TEXT
+);
+
+insertTest:
+INSERT INTO test (id2) VALUES(?);
+
+insertTest2:
+INSERT INTO test2 (other) VALUES(?);
+
+updateTestId:
+UPDATE test
+SET id = t2.id2
+FROM (
+  SELECT
+    id2
+  FROM test2
+) AS t2 RETURNING test.id;
+
+updateTestId2:
+UPDATE test
+SET id2 = t2.id2
+FROM test2 t2
+WHERE other = ? RETURNING test.id2;
diff --git a/...n-postgresql/src/test/kotlin/app/cash/sqldelight/postgresql/integration/PostgreSqlTest.kt b/...n-postgresql/src/test/kotlin/app/cash/sqldelight/postgresql/integration/PostgreSqlTest.kt
@@ -783,6 +783,24 @@ class PostgreSqlTest {
     }
   }
 
+  @Test
+  fun testUpdateSetFromId() {
+    database.updatesQueries.insertTest(31)
+    database.updatesQueries.insertTest2("X")
+    with(database.updatesQueries.updateTestId().executeAsOne()) {
+      assertThat(this).isEqualTo(1)
+    }
+  }
+
+  @Test
+  fun testUpdateSetFromId2() {
+    database.updatesQueries.insertTest(31)
+    database.updatesQueries.insertTest2("X")
+    with(database.updatesQueries.updateTestId2("X").executeAsOne()) {
+      assertThat(id2).isEqualTo(1)
+    }
+  }
+
   @Test
   fun testSelectTsVectorSearch() {
     database.textSearchQueries.insertLiteral("the rain in spain")