sqldelight · AlecKazakova · Feb 15, 2024 · Feb 13, 2024 · Feb 13, 2024 · Feb 13, 2024
diff --git a/...postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/PostgreSql.bnf b/...postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/PostgreSql.bnf
@@ -314,7 +314,9 @@ alter_table_drop_column ::= DROP [ COLUMN ] {column_name} {
   pin = 1
 }
 
-alter_table_add_constraint ::= ADD table_constraint
+alter_table_add_constraint ::= ADD table_constraint {
+    mixin = "app.cash.sqldelight.dialects.postgresql.grammar.mixins.AlterTableAddConstraintMixin"
+}
 
 type_clause ::= 'TYPE'
 

diff --git a/...in/app/cash/sqldelight/dialects/postgresql/grammar/mixins/AlterTableAddConstraintMixin.kt b/...in/app/cash/sqldelight/dialects/postgresql/grammar/mixins/AlterTableAddConstraintMixin.kt
@@ -0,0 +1,27 @@
+package app.cash.sqldelight.dialects.postgresql.grammar.mixins
+
+import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlAlterTableAddConstraint
+import com.alecstrong.sql.psi.core.psi.AlterTableApplier
+import com.alecstrong.sql.psi.core.psi.LazyQuery
+import com.alecstrong.sql.psi.core.psi.SqlCompositeElementImpl
+import com.alecstrong.sql.psi.core.psi.SqlTypes
+import com.intellij.lang.ASTNode
+
+abstract class AlterTableAddConstraintMixin(node: ASTNode) :
+  SqlCompositeElementImpl(node),
+  PostgreSqlAlterTableAddConstraint,
+  AlterTableApplier {
+  override fun applyTo(lazyQuery: LazyQuery): LazyQuery =
+    if (tableConstraint.node.findChildByType(SqlTypes.PRIMARY) != null &&
+      tableConstraint.node.findChildByType(SqlTypes.KEY) != null
+    ) {
+      val columns = lazyQuery.query.columns.map { queryCol ->
+        tableConstraint.indexedColumnList.find { indexedCol -> queryCol.element.textMatches(indexedCol) }.let {
+          queryCol.copy(nullable = if (it != null) false else queryCol.nullable)
+        }
+      }
+      LazyQuery(lazyQuery.tableName, query = { lazyQuery.query.copy(columns = columns) })
+    } else {
+      lazyQuery
+    }
+}
diff --git a/...elight-compiler/src/test/kotlin/app/cash/sqldelight/core/migrations/MigrationQueryTest.kt b/...elight-compiler/src/test/kotlin/app/cash/sqldelight/core/migrations/MigrationQueryTest.kt
@@ -29,6 +29,10 @@ class MigrationQueryTest {
     checkFixtureCompiles("alter-table-alter-column", PostgreSqlDialect())
   }
 
+  @Test fun `alter table add constraint`() {
+    checkFixtureCompiles("alter-table-add-constraint", PostgreSqlDialect())
+  }
+
   @Test fun `varying query migration packages`() {
     checkFixtureCompiles("varying-query-migration-packages", PostgreSqlDialect())
   }

diff --git a/...mpiler/src/test/migration-interface-fixtures/alter-table-add-constraint/com/example/1.sqm b/...mpiler/src/test/migration-interface-fixtures/alter-table-add-constraint/com/example/1.sqm
@@ -0,0 +1,9 @@
+CREATE TABLE TestSingle(
+  first INTEGER,
+  second TEXT
+);
+
+CREATE TABLE TestCompound (
+  first INTEGER,
+  second TEXT
+);
diff --git a/...mpiler/src/test/migration-interface-fixtures/alter-table-add-constraint/com/example/2.sqm b/...mpiler/src/test/migration-interface-fixtures/alter-table-add-constraint/com/example/2.sqm
@@ -0,0 +1,3 @@
+ALTER TABLE TestSingle ADD PRIMARY KEY (first);
+ALTER TABLE TestCompound ADD CONSTRAINT pk_first_second PRIMARY KEY (first, second);
+
diff --git a/...iler/src/test/migration-interface-fixtures/alter-table-add-constraint/com/example/Data.sq b/...iler/src/test/migration-interface-fixtures/alter-table-add-constraint/com/example/Data.sq
@@ -0,0 +1,8 @@
+selectSingle:
+SELECT *
+FROM TestSingle;
+
+selectCompound:
+SELECT *
+FROM TestCompound;
+
diff --git a/...migration-interface-fixtures/alter-table-add-constraint/output/com/example/DataQueries.kt b/...migration-interface-fixtures/alter-table-add-constraint/output/com/example/DataQueries.kt
@@ -0,0 +1,51 @@
+package com.example
+
+import app.cash.sqldelight.Query
+import app.cash.sqldelight.TransacterImpl
+import app.cash.sqldelight.db.SqlDriver
+import app.cash.sqldelight.driver.jdbc.JdbcCursor
+import kotlin.Any
+import kotlin.Int
+import kotlin.String
+
+public class DataQueries(
+  driver: SqlDriver,
+) : TransacterImpl(driver) {
+  public fun <T : Any> selectSingle(mapper: (first: Int, second: String?) -> T): Query<T> =
+      Query(-79_317_191, arrayOf("TestSingle"), driver, "Data.sq", "selectSingle", """
+  |SELECT *
+  |FROM TestSingle
+  """.trimMargin()) { cursor ->
+    check(cursor is JdbcCursor)
+    mapper(
+      cursor.getInt(0)!!,
+      cursor.getString(1)
+    )
+  }
+
+  public fun selectSingle(): Query<TestSingle> = selectSingle { first, second ->
+    TestSingle(
+      first,
+      second
+    )
+  }
+
+  public fun <T : Any> selectCompound(mapper: (first: Int, second: String) -> T): Query<T> =
+      Query(-19_725_220, arrayOf("TestCompound"), driver, "Data.sq", "selectCompound", """
+  |SELECT *
+  |FROM TestCompound
+  """.trimMargin()) { cursor ->
+    check(cursor is JdbcCursor)
+    mapper(
+      cursor.getInt(0)!!,
+      cursor.getString(1)!!
+    )
+  }
+
+  public fun selectCompound(): Query<TestCompound> = selectCompound { first, second ->
+    TestCompound(
+      first,
+      second
+    )
+  }
+}
diff --git a/...igration-interface-fixtures/alter-table-add-constraint/output/com/example/TestCompound.kt b/...igration-interface-fixtures/alter-table-add-constraint/output/com/example/TestCompound.kt
@@ -0,0 +1,9 @@
+package com.example
+
+import kotlin.Int
+import kotlin.String
+
+public data class TestCompound(
+  public val first: Int,
+  public val second: String,
+)
diff --git a/.../migration-interface-fixtures/alter-table-add-constraint/output/com/example/TestSingle.kt b/.../migration-interface-fixtures/alter-table-add-constraint/output/com/example/TestSingle.kt
@@ -0,0 +1,9 @@
+package com.example
+
+import kotlin.Int
+import kotlin.String
+
+public data class TestSingle(
+  public val first: Int,
+  public val second: String?,
+)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		ALTER TABLE TestSingle ADD PRIMARY KEY (first);
		ALTER TABLE TestCompound ADD CONSTRAINT pk_first_second PRIMARY KEY (first, second);