[KYUUBI #7708][SERVER] Reconcile zero-count metadata updates

kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala

@@ -337,16 +337,22 @@ class MetadataManager extends AbstractService("MetadataManager") {
                     info(s"Retrying metadata requests for $id")
                     var request = ref.metadataRequests.peek()
                     while (request != null) {
-                      request match {
-                        case insert: InsertMetadata =>
-                          insertMetadata(insert.metadata, asyncRetryOnError = false)
-                        case update: UpdateMetadata =>
-                          updateMetadata(update.metadata, asyncRetryOnError = false)
+                      try {
+                        request match {
+                          case insert: InsertMetadata =>
+                            insertMetadata(insert.metadata, asyncRetryOnError = false)
+                          case update: UpdateMetadata =>
+                            updateMetadata(update.metadata, asyncRetryOnError = false)
+                        }
+                        removeRetryRequest(ref, request)
+                      } catch {
+                        case e: MetadataUpdatePostconditionException =>
+                          error(
+                            s"Discarding metadata update retry for $id after " +
+                              "its postcondition check failed",
+                            e)
+                          removeRetryRequest(ref, request)
                       }
-                      ref.metadataRequests.remove(request)
-                      MetricsSystem.tracing(_.markMeter(
-                        MetricsConstants.METADATA_REQUEST_RETRYING,
-                        -1L))
                       request = ref.metadataRequests.peek()
                     }
                   } catch {
@@ -379,6 +385,13 @@ class MetadataManager extends AbstractService("MetadataManager") {
       requestsRetryInterval,
       TimeUnit.MILLISECONDS)
   }
+
+  private def removeRetryRequest(
+      ref: MetadataRequestsRetryRef,
+      request: MetadataRequest): Unit = {
+    ref.metadataRequests.remove(request)
+    MetricsSystem.tracing(_.markMeter(MetricsConstants.METADATA_REQUEST_RETRYING, -1L))
+  }
 }
 
 object MetadataManager extends Logging {

kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala

@@ -19,6 +19,7 @@ package org.apache.kyuubi.server.metadata
 
 import java.io.Closeable
 
+import org.apache.kyuubi.KyuubiException
 import org.apache.kyuubi.server.metadata.api.{KubernetesEngineInfo, Metadata, MetadataFilter}
 
 trait MetadataStore extends Closeable {
@@ -118,3 +119,21 @@ trait MetadataStore extends Closeable {
    */
   def cleanupKubernetesEngineInfoByAge(maxAge: Long, limit: Int): Int
 }
+
+sealed abstract private[metadata] class MetadataUpdatePostconditionException(message: String)
+  extends KyuubiException(message)
+
+private[metadata] class MetadataRowNotFoundException(identifier: String)
+  extends MetadataUpdatePostconditionException(
+    s"Metadata row $identifier was not found after an update returned 0")
+
+private[metadata] class MetadataUpdateMismatchException(
+    identifier: String,
+    mismatchedValues: Seq[(String, String, String)])
+  extends MetadataUpdatePostconditionException(
+    s"Metadata row $identifier does not match the requested update for columns: " +
+      mismatchedValues.map { case (column, expected, actual) =>
+        s"$column (expected: $expected, actual: $actual)"
+      }.mkString(", ")) {
+  val mismatchedColumns: Seq[String] = mismatchedValues.map(_._1)
+}

kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala

@@ -34,7 +34,7 @@ import com.zaxxer.hikari.{HikariConfig, HikariDataSource}
 import org.apache.kyuubi.{KyuubiException, Logging, Utils}
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.operation.OperationState
-import org.apache.kyuubi.server.metadata.MetadataStore
+import org.apache.kyuubi.server.metadata.{MetadataRowNotFoundException, MetadataStore, MetadataUpdateMismatchException}
 import org.apache.kyuubi.server.metadata.api.{KubernetesEngineInfo, Metadata, MetadataFilter}
 import org.apache.kyuubi.server.metadata.jdbc.DatabaseType._
 import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._
@@ -332,76 +332,106 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
 
   override def updateMetadata(metadata: Metadata): Unit = {
     val queryBuilder = new StringBuilder
-    val params = ListBuffer[Any]()
+    val updateFields = ListBuffer[(String, Any)]()
 
     queryBuilder.append(s"UPDATE $METADATA_TABLE")
-    val setClauses = ListBuffer[String]()
     Option(metadata.kyuubiInstance).foreach { _ =>
-      setClauses += "kyuubi_instance = ?"
-      params += metadata.kyuubiInstance
+      updateFields += (("kyuubi_instance", metadata.kyuubiInstance))
     }
     Option(metadata.state).foreach { _ =>
-      setClauses += "state = ?"
-      params += metadata.state
+      updateFields += (("state", metadata.state))
     }
     Option(metadata.requestConf).filter(_.nonEmpty).foreach { _ =>
-      setClauses += "request_conf =?"
-      params += valueAsString(metadata.requestConf)
+      updateFields += (("request_conf", valueAsString(metadata.requestConf)))
     }
     metadata.clusterManager.foreach { cm =>
-      setClauses += "cluster_manager = ?"
-      params += cm
+      updateFields += (("cluster_manager", cm))
     }
     if (metadata.endTime > 0) {
-      setClauses += "end_time = ?"
-      params += metadata.endTime
+      updateFields += (("end_time", metadata.endTime))
     }
     if (metadata.engineOpenTime > 0) {
-      setClauses += "engine_open_time = ?"
-      params += metadata.engineOpenTime
+      updateFields += (("engine_open_time", metadata.engineOpenTime))
     }
     Option(metadata.engineId).foreach { _ =>
-      setClauses += "engine_id = ?"
-      params += metadata.engineId
+      updateFields += (("engine_id", metadata.engineId))
     }
     Option(metadata.engineName).foreach { _ =>
-      setClauses += "engine_name = ?"
-      params += metadata.engineName
+      updateFields += (("engine_name", metadata.engineName))
     }
     Option(metadata.engineUrl).foreach { _ =>
-      setClauses += "engine_url = ?"
-      params += metadata.engineUrl
+      updateFields += (("engine_url", metadata.engineUrl))
     }
     Option(metadata.engineState).foreach { _ =>
-      setClauses += "engine_state = ?"
-      params += metadata.engineState
+      updateFields += (("engine_state", metadata.engineState))
     }
     metadata.engineError.foreach { error =>
-      setClauses += "engine_error = ?"
-      params += error
+      updateFields += (("engine_error", error))
     }
     if (metadata.peerInstanceClosed) {
-      setClauses += "peer_instance_closed = ?"
-      params += metadata.peerInstanceClosed
+      updateFields += (("peer_instance_closed", metadata.peerInstanceClosed))
     }
-    if (setClauses.nonEmpty) {
-      queryBuilder.append(setClauses.mkString(" SET ", ", ", ""))
+    if (updateFields.nonEmpty) {
+      queryBuilder.append(updateFields.map { case (column, _) =>
+        s"$column = ?"
+      }.mkString(" SET ", ", ", ""))
     }
     queryBuilder.append(" WHERE identifier = ?")
-    params += metadata.identifier
+    val params = updateFields.map(_._2) :+ metadata.identifier
 
     val query = queryBuilder.toString()
     JdbcUtils.withConnection { connection =>
       withUpdateCount(connection, query, params.toSeq: _*) { updateCount =>
         if (updateCount == 0) {
-          throw new KyuubiException(
-            s"Error updating metadata for ${metadata.identifier} by SQL: $query, " +
-              s"with params: ${params.mkString(", ")}")
+          verifyMetadataUpdatePostcondition(connection, metadata.identifier, updateFields.toSeq)
+        }
+      }
+    }
+  }
+
+  private def verifyMetadataUpdatePostcondition(
+      connection: Connection,
+      identifier: String,
+      updateFields: Seq[(String, Any)]): Unit = {
+    val query = s"SELECT ${updateFields.map(_._1).mkString(", ")} " +
+      s"FROM $METADATA_TABLE WHERE identifier = ?"
+    withResultSet(connection, query, identifier) { resultSet =>
+      if (!resultSet.next()) {
+        throw new MetadataRowNotFoundException(identifier)
+      }
+
+      val mismatchedValues = updateFields.zipWithIndex.flatMap { case ((column, expected), index) =>
+        val actual = getResultSetValue(resultSet, index + 1, expected)
+        if (actual == expected) {
+          None
+        } else {
+          val Seq((_, redactedExpected), (_, redactedActual)) = Utils.redact(
+            Some(conf.get(KyuubiConf.SERVER_SECRET_REDACTION_PATTERN)),
+            Seq(column -> String.valueOf(expected), column -> String.valueOf(actual)))
+          Some((column, redactedExpected, redactedActual))
         }
       }
+      if (mismatchedValues.nonEmpty) {
+        throw new MetadataUpdateMismatchException(identifier, mismatchedValues)
+      }
+      warn(s"Update for $identifier returned 0 but all requested values already match, " +
+        "treating as idempotent success")
     }
   }
 
+  private def getResultSetValue(resultSet: ResultSet, index: Int, expected: Any): Any = {
+    val value = expected match {
+      case _: String => resultSet.getString(index)
+      case _: Int => resultSet.getInt(index)
+      case _: Long => resultSet.getLong(index)
+      case _: Double => resultSet.getDouble(index)
+      case _: Float => resultSet.getFloat(index)
+      case _: Boolean => resultSet.getBoolean(index)
+      case _ => resultSet.getObject(index)
+    }
+    if (resultSet.wasNull()) null else value
+  }
+
   override def cleanupMetadataByIdentifier(identifier: String): Unit = {
     val query = s"DELETE FROM $METADATA_TABLE WHERE identifier = ?"
     JdbcUtils.withConnection { connection =>

kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/MetadataManagerSuite.scala

@@ -97,6 +97,22 @@ class MetadataManagerSuite extends KyuubiFunSuite {
     }
   }
 
+  test("stop retrying an update when its row remains missing") {
+    withMetadataManager(Map(
+      METADATA_REQUEST_ASYNC_RETRY_ENABLED.key -> "true",
+      METADATA_REQUEST_RETRY_INTERVAL.key -> "100")) { metadataManager =>
+      val metadata = newMetadata()
+      metadataManager.updateMetadata(metadata)
+      val retryRef = metadataManager.getMetadataRequestsRetryRef(metadata.identifier)
+
+      assert(retryRef != null)
+      eventually(timeout(3.seconds)) {
+        assert(!retryRef.hasRemainingRequests())
+        assert(metadataManager.getMetadataRequestsRetryRef(metadata.identifier) == null)
+      }
+    }
+  }
+
   test("async metadata request metrics") {
     withMetadataManager(Map(
       METADATA_REQUEST_ASYNC_RETRY_ENABLED.key -> "true",

kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala

@@ -22,13 +22,14 @@ import java.util.UUID
 import org.scalatest.concurrent.PatienceConfiguration.Timeout
 import org.scalatest.time.SpanSugar._
 
-import org.apache.kyuubi.{KyuubiException, KyuubiFunSuite}
+import org.apache.kyuubi.KyuubiFunSuite
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.engine.ApplicationState
-import org.apache.kyuubi.server.metadata.MetadataManager
+import org.apache.kyuubi.server.metadata.{MetadataManager, MetadataRowNotFoundException, MetadataUpdateMismatchException}
 import org.apache.kyuubi.server.metadata.api.{KubernetesEngineInfo, Metadata, MetadataFilter}
 import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._
 import org.apache.kyuubi.session.SessionType
+import org.apache.kyuubi.util.command.CommandLineUtils.REDACTION_REPLACEMENT_TEXT
 
 class JDBCMetadataStoreSuite extends KyuubiFunSuite {
   private val conf = KyuubiConf()
@@ -274,13 +275,102 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
     jdbcMetadataStore.cleanupMetadataByIdentifier(batchId)
   }
 
-  test("throw exception if update count is 0") {
+  test("report a missing metadata row after update count is 0") {
     val metadata = Metadata(identifier = UUID.randomUUID().toString, state = "RUNNING")
-    intercept[KyuubiException] {
+    intercept[MetadataRowNotFoundException] {
       jdbcMetadataStore.updateMetadata(metadata)
     }
   }
 
+  test("accept update count 0 if the metadata already matches") {
+    val batchId = UUID.randomUUID().toString
+    val triggerName = s"ignore_metadata_update_${batchId.replace("-", "")}"
+    val metadata = Metadata(
+      identifier = batchId,
+      sessionType = SessionType.BATCH,
+      realUser = "kyuubi",
+      username = "kyuubi",
+      state = "RUNNING",
+      createTime = System.currentTimeMillis(),
+      engineType = "spark")
+    val metadataToUpdate = Metadata(
+      identifier = batchId,
+      kyuubiInstance = "localhost:10099",
+      state = "RUNNING",
+      requestConf = Map("spark.master" -> "local", "jdbc.password" -> "stored-credential"),
+      clusterManager = Some("kubernetes"),
+      engineOpenTime = System.currentTimeMillis(),
+      engineId = "app_id",
+      engineName = "app_name",
+      engineUrl = "app_url?token=stored-credential",
+      engineState = "FAILED",
+      engineError = Some("engine_error"),
+      endTime = System.currentTimeMillis(),
+      peerInstanceClosed = true)
+
+    try {
+      jdbcMetadataStore.insertMetadata(metadata)
+      jdbcMetadataStore.updateMetadata(metadataToUpdate)
+      executeSql(
+        s"""
+           |CREATE TRIGGER $triggerName
+           |BEFORE UPDATE ON metadata
+           |WHEN OLD.identifier = '$batchId'
+           |BEGIN
+           |  SELECT RAISE(IGNORE);
+           |END
+           |""".stripMargin)
+
+      jdbcMetadataStore.updateMetadata(metadataToUpdate)
+
+      val error = intercept[MetadataUpdateMismatchException] {
+        jdbcMetadataStore.updateMetadata(metadataToUpdate.copy(state = "FINISHED"))
+      }
+      assert(error.mismatchedColumns === Seq("state"))
+      assert(error.getMessage.contains("state (expected: FINISHED, actual: RUNNING)"))
+
+      val timeError = intercept[MetadataUpdateMismatchException] {
+        jdbcMetadataStore.updateMetadata(
+          metadataToUpdate.copy(engineOpenTime = metadataToUpdate.engineOpenTime + 1))
+      }
+      assert(timeError.getMessage.contains(
+        s"engine_open_time (expected: ${metadataToUpdate.engineOpenTime + 1}, " +
+          s"actual: ${metadataToUpdate.engineOpenTime})"))
+
+      val sensitiveError = intercept[MetadataUpdateMismatchException] {
+        jdbcMetadataStore.updateMetadata(metadataToUpdate.copy(
+          requestConf = Map("jdbc.password" -> "requested-credential"),
+          engineUrl = "app_url?token=requested-credential"))
+      }
+      assert(sensitiveError.mismatchedColumns === Seq("request_conf", "engine_url"))
+      Seq("request_conf", "engine_url").foreach { column =>
+        assert(sensitiveError.getMessage.contains(
+          s"$column (expected: $REDACTION_REPLACEMENT_TEXT, actual: $REDACTION_REPLACEMENT_TEXT)"))
+      }
+      assert(!sensitiveError.getMessage.contains("stored-credential"))
+      assert(!sensitiveError.getMessage.contains("requested-credential"))
+
+      val customRedactionStore = new JDBCMetadataStore(conf.clone
+        .set(METADATA_STORE_JDBC_DATABASE_SCHEMA_INIT, false)
+        .set(KyuubiConf.SERVER_SECRET_REDACTION_PATTERN, "app_name|private".r))
+      try {
+        val customError = intercept[MetadataUpdateMismatchException] {
+          customRedactionStore.updateMetadata(metadataToUpdate.copy(engineName = "private-name"))
+        }
+        assert(customError.getMessage.contains(
+          s"engine_name (expected: $REDACTION_REPLACEMENT_TEXT, " +
+            s"actual: $REDACTION_REPLACEMENT_TEXT)"))
+        assert(!customError.getMessage.contains("app_name"))
+        assert(!customError.getMessage.contains("private-name"))
+      } finally {
+        customRedactionStore.close()
+      }
+    } finally {
+      executeSql(s"DROP TRIGGER IF EXISTS $triggerName")
+      jdbcMetadataStore.cleanupMetadataByIdentifier(batchId)
+    }
+  }
+
   test("get schema urls with correct version ordering") {
     val url1 = "metadata-store-schema-1.7.0.mysql.sql"
     val url2 = "metadata-store-schema-1.7.1.mysql.sql"
@@ -363,4 +453,18 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
     jdbcMetadataStore.cleanupKubernetesEngineInfoByIdentifier(tag)
     assert(jdbcMetadataStore.getKubernetesMetaEngineInfo(tag) == null)
   }
+
+  private def executeSql(sql: String): Unit = {
+    val connection = jdbcMetadataStore.hikariDataSource.getConnection
+    try {
+      val statement = connection.createStatement()
+      try {
+        statement.execute(sql)
+      } finally {
+        statement.close()
+      }
+    } finally {
+      connection.close()
+    }
+  }
 }