diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala index 18b840ddb..d7e2f06cf 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala @@ -269,9 +269,14 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging { queryBuilder.append(" WHERE identifier = ? ") params += metadata.identifier + val query = queryBuilder.toString() withConnection() { connection => - execute(connection, queryBuilder.toString(), params: _*) - + withUpdateCount(connection, query, params: _*) { updateCount => + if (updateCount == 0) { + throw new KyuubiException( + s"Error updating metadata for ${metadata.identifier} with $query") + } + } } } @@ -397,6 +402,26 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging { } } + private def withUpdateCount[T]( + conn: Connection, + sql: String, + params: Any*)(f: Int => T): T = { + debug(s"executing sql $sql with update count") + var statement: PreparedStatement = null + try { + statement = conn.prepareStatement(sql) + setStatementParams(statement, params: _*) + f(statement.executeUpdate()) + } catch { + case e: SQLException => + throw new KyuubiException(e.getMessage, e) + } finally { + if (statement != null) { + Utils.tryLogNonFatalError(statement.close()) + } + } + } + private def setStatementParams(statement: PreparedStatement, params: Any*): Unit = { params.zipWithIndex.foreach { case (param, index) => param match { diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala index c7adf4a2c..106b9a0b4 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala @@ -22,7 +22,7 @@ import java.util.UUID import org.scalatest.concurrent.PatienceConfiguration.Timeout import org.scalatest.time.SpanSugar._ -import org.apache.kyuubi.KyuubiFunSuite +import org.apache.kyuubi.{KyuubiException, KyuubiFunSuite} import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter} import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._ @@ -266,4 +266,11 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite { assert(jdbcMetadataStore.getMetadata(batchId, true) == null) } } + + test("throw exception if update count is 0") { + val metadata = Metadata(identifier = UUID.randomUUID().toString, state = "RUNNING") + intercept[KyuubiException] { + jdbcMetadataStore.updateMetadata(metadata) + } + } }