Validate the optimistic lock in INSERT ON CONFLICT DO UPDATE

CHANGELOG.md

@@ -25,6 +25,7 @@
 
 ### Fixed
 - [Compiler] Suppress Kotlin extra warnings in generated code (#6208 by @eyupcanakman)
+- [SQLite Dialect] Validate optimistic lock column and WHERE clause in INSERT ... ON CONFLICT DO UPDATE (#6306 by @eyupcanakman)
 - [Compiler] Other columns in a non-grouped aggregate result set are always nullable
 - [PostgreSQL Dialect] Resolve nullability correctly for coalesce and ifnull
 - [PostgreSQL Dialect] Fixed IDE integration of the PostgreSQL dialect

dialects/sqlite-3-24/src/main/kotlin/app/cash/sqldelight/dialects/sqlite_3_24/grammar/mixins/InsertStmtMixin.kt

@@ -1,10 +1,22 @@
 package app.cash.sqldelight.dialects.sqlite_3_24.grammar.mixins
 
 import app.cash.sqldelight.dialects.sqlite_3_24.grammar.psi.SqliteInsertStmt
+import app.cash.sqldelight.dialects.sqlite_3_24.grammar.psi.SqliteUpsertDoUpdate
 import com.alecstrong.sql.psi.core.SqlAnnotationHolder
+import com.alecstrong.sql.psi.core.psi.SqlBinaryAddExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryEqualityExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryExpr
+import com.alecstrong.sql.psi.core.psi.SqlBindExpr
+import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
+import com.alecstrong.sql.psi.core.psi.SqlColumnName
+import com.alecstrong.sql.psi.core.psi.SqlExpr
+import com.alecstrong.sql.psi.core.psi.SqlSetterExpression
 import com.alecstrong.sql.psi.core.psi.SqlTypes
+import com.alecstrong.sql.psi.core.psi.SqlUpdateStmtSubsequentSetter
 import com.alecstrong.sql.psi.core.psi.impl.SqlInsertStmtImpl
+import com.alecstrong.sql.psi.core.psi.mixins.ColumnDefMixin
 import com.intellij.lang.ASTNode
+import com.intellij.psi.util.PsiTreeUtil
 
 internal abstract class InsertStmtMixin(
   node: ASTNode,
@@ -50,6 +62,91 @@ internal abstract class InsertStmtMixin(
             "also specifying a conflict resolution algorithm ($conflictResolution)",
         )
       }
+
+      if (upsertDoUpdate != null) {
+        validateOptimisticLock(upsertDoUpdate, annotationHolder)
+      }
+    }
+  }
+
+  private fun validateOptimisticLock(
+    doUpdate: SqliteUpsertDoUpdate,
+    annotationHolder: SqlAnnotationHolder,
+  ) {
+    val table = tablesAvailable(this)
+      .firstOrNull { it.tableName.name == tableName.name } ?: return
+    val lock = table.query.columns
+      .mapNotNull { it.element.parent as? ColumnDefMixin }
+      .singleOrNull { col ->
+        col.columnType.node.getChildren(null).any { it.text == "LOCK" }
+      } ?: return
+
+    val firstColumnName = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlColumnName::class.java).firstOrNull()
+    val firstSetterExpr = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlSetterExpression::class.java).firstOrNull()
+    val subsequentSetters = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlUpdateStmtSubsequentSetter::class.java)
+
+    val columns = listOfNotNull(firstColumnName) + subsequentSetters.mapNotNull { it.columnName }
+    val setters = listOfNotNull(firstSetterExpr) + subsequentSetters.mapNotNull { it.setterExpression }
+
+    val setter = columns.zip(setters)
+      .singleOrNull { (col, _) -> col.textMatches(lock.columnName.name) }
+      ?.second
+
+    if (setter == null) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        "This statement is missing the optimistic lock in its SET clause.",
+      )
+      return
+    }
+
+    // excluded holds the row being inserted, so only the target table carries the stored lock.
+    val storedNames = listOfNotNull(tableName.name, tableAlias?.name)
+    fun SqlColumnExpr.isStoredLock() = columnName.textMatches(lock.columnName.name) &&
+      (tableName?.let { qualifier -> qualifier.name in storedNames } ?: true)
+
+    val setterExpressions = (setter.expr as? SqlBinaryExpr)?.getExprList()
+    val selfIncrementsByOne = setter.expr is SqlBinaryAddExpr &&
+      setter.expr.node.getChildren(null).any { it.text == "+" } &&
+      setterExpressions?.getOrNull(1)?.textMatches("1") ?: false
+    val bindIncrement = selfIncrementsByOne &&
+      setterExpressions.getOrNull(0) is SqlBindExpr
+    val selfIncrements = selfIncrementsByOne &&
+      (setterExpressions.getOrNull(0) as? SqlColumnExpr)?.isStoredLock() == true
+
+    // Confirm the statement has SET lock = :arg + 1 or SET lock = lock + 1
+    if (!bindIncrement && !selfIncrements) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        """The optimistic lock must be set exactly like "${lock.columnName.name} = :${lock.columnName.name} + 1".""",
+      )
+      return
+    }
+
+    // Unlike a multirow UPDATE, DO UPDATE only touches the row that conflicted, so it still needs the WHERE clause.
+    val whereExpr = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlExpr::class.java).firstOrNull()
+    if (whereExpr == null) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        "This statement is missing a WHERE clause to check the optimistic lock.",
+      )
+      return
+    }
+
+    val equalityExprs = buildList<SqlBinaryEqualityExpr> {
+      if (whereExpr is SqlBinaryEqualityExpr) add(whereExpr)
+      addAll(PsiTreeUtil.findChildrenOfType(whereExpr, SqlBinaryEqualityExpr::class.java))
+    }
+    if (equalityExprs.none { eq ->
+        eq.node.getChildren(null).any { it.text == "=" || it.text == "==" } &&
+          (eq.getExprList()[0] as? SqlColumnExpr)?.isStoredLock() == true &&
+          eq.getExprList()[1] is SqlBindExpr
+      }
+    ) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        """The optimistic lock must be queried exactly like "${lock.columnName.name} == :${lock.columnName.name}".""",
+      )
     }
   }
 }

dialects/sqlite-3-35/src/main/kotlin/app/cash/sqldelight/dialects/sqlite_3_35/grammar/mixins/InsertStmtMixin.kt

@@ -1,10 +1,22 @@
 package app.cash.sqldelight.dialects.sqlite_3_35.grammar.mixins
 
 import app.cash.sqldelight.dialects.sqlite_3_35.grammar.psi.SqliteInsertStmt
+import app.cash.sqldelight.dialects.sqlite_3_35.grammar.psi.SqliteUpsertDoUpdate
 import com.alecstrong.sql.psi.core.SqlAnnotationHolder
+import com.alecstrong.sql.psi.core.psi.SqlBinaryAddExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryEqualityExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryExpr
+import com.alecstrong.sql.psi.core.psi.SqlBindExpr
+import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
+import com.alecstrong.sql.psi.core.psi.SqlColumnName
+import com.alecstrong.sql.psi.core.psi.SqlExpr
+import com.alecstrong.sql.psi.core.psi.SqlSetterExpression
 import com.alecstrong.sql.psi.core.psi.SqlTypes
+import com.alecstrong.sql.psi.core.psi.SqlUpdateStmtSubsequentSetter
 import com.alecstrong.sql.psi.core.psi.impl.SqlInsertStmtImpl
+import com.alecstrong.sql.psi.core.psi.mixins.ColumnDefMixin
 import com.intellij.lang.ASTNode
+import com.intellij.psi.util.PsiTreeUtil
 
 internal abstract class InsertStmtMixin(
   node: ASTNode,
@@ -62,6 +74,91 @@ internal abstract class InsertStmtMixin(
           "Only the final ON CONFLICT clause may omit the conflict target",
         )
       }
+
+      if (upsertDoUpdate != null) {
+        validateOptimisticLock(upsertDoUpdate, annotationHolder)
+      }
+    }
+  }
+
+  private fun validateOptimisticLock(
+    doUpdate: SqliteUpsertDoUpdate,
+    annotationHolder: SqlAnnotationHolder,
+  ) {
+    val table = tablesAvailable(this)
+      .firstOrNull { it.tableName.name == tableName.name } ?: return
+    val lock = table.query.columns
+      .mapNotNull { it.element.parent as? ColumnDefMixin }
+      .singleOrNull { col ->
+        col.columnType.node.getChildren(null).any { it.text == "LOCK" }
+      } ?: return
+
+    val firstColumnName = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlColumnName::class.java).firstOrNull()
+    val firstSetterExpr = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlSetterExpression::class.java).firstOrNull()
+    val subsequentSetters = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlUpdateStmtSubsequentSetter::class.java)
+
+    val columns = listOfNotNull(firstColumnName) + subsequentSetters.mapNotNull { it.columnName }
+    val setters = listOfNotNull(firstSetterExpr) + subsequentSetters.mapNotNull { it.setterExpression }
+
+    val setter = columns.zip(setters)
+      .singleOrNull { (col, _) -> col.textMatches(lock.columnName.name) }
+      ?.second
+
+    if (setter == null) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        "This statement is missing the optimistic lock in its SET clause.",
+      )
+      return
+    }
+
+    // excluded holds the row being inserted, so only the target table carries the stored lock.
+    val storedNames = listOfNotNull(tableName.name, tableAlias?.name)
+    fun SqlColumnExpr.isStoredLock() = columnName.textMatches(lock.columnName.name) &&
+      (tableName?.let { qualifier -> qualifier.name in storedNames } ?: true)
+
+    val setterExpressions = (setter.expr as? SqlBinaryExpr)?.getExprList()
+    val selfIncrementsByOne = setter.expr is SqlBinaryAddExpr &&
+      setter.expr.node.getChildren(null).any { it.text == "+" } &&
+      setterExpressions?.getOrNull(1)?.textMatches("1") ?: false
+    val bindIncrement = selfIncrementsByOne &&
+      setterExpressions.getOrNull(0) is SqlBindExpr
+    val selfIncrements = selfIncrementsByOne &&
+      (setterExpressions.getOrNull(0) as? SqlColumnExpr)?.isStoredLock() == true
+
+    // Confirm the statement has SET lock = :arg + 1 or SET lock = lock + 1
+    if (!bindIncrement && !selfIncrements) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        """The optimistic lock must be set exactly like "${lock.columnName.name} = :${lock.columnName.name} + 1".""",
+      )
+      return
+    }
+
+    // Unlike a multirow UPDATE, DO UPDATE only touches the row that conflicted, so it still needs the WHERE clause.
+    val whereExpr = PsiTreeUtil.getChildrenOfTypeAsList(doUpdate, SqlExpr::class.java).firstOrNull()
+    if (whereExpr == null) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        "This statement is missing a WHERE clause to check the optimistic lock.",
+      )
+      return
+    }
+
+    val equalityExprs = buildList<SqlBinaryEqualityExpr> {
+      if (whereExpr is SqlBinaryEqualityExpr) add(whereExpr)
+      addAll(PsiTreeUtil.findChildrenOfType(whereExpr, SqlBinaryEqualityExpr::class.java))
+    }
+    if (equalityExprs.none { eq ->
+        eq.node.getChildren(null).any { it.text == "=" || it.text == "==" } &&
+          (eq.getExprList()[0] as? SqlColumnExpr)?.isStoredLock() == true &&
+          eq.getExprList()[1] is SqlBindExpr
+      }
+    ) {
+      annotationHolder.createErrorAnnotation(
+        doUpdate,
+        """The optimistic lock must be queried exactly like "${lock.columnName.name} == :${lock.columnName.name}".""",
+      )
     }
   }
 }