Save

Author: turboFei

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

@@ -17,11 +17,13 @@
 
 package org.apache.kyuubi.server.metadata
 
-import com.google.common.annotations.VisibleForTesting
-
 import java.util.concurrent.{ConcurrentHashMap, ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.atomic.AtomicInteger
+
 import scala.collection.JavaConverters._
+
+import com.google.common.annotations.VisibleForTesting
+
 import org.apache.kyuubi.{KyuubiException, Logging}
 import org.apache.kyuubi.client.api.v1.dto.Batch
 import org.apache.kyuubi.config.KyuubiConf

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

@@ -106,6 +106,10 @@ trait MetadataStore extends Closeable {
    */
   def cleanupMetadataByAge(maxAge: Long, limit: Int): Int
 
+  def cleanupMetadataByAge(maxAge: Long): Int = {
+    cleanupMetadataByAge(maxAge, Int.MaxValue)
+  }
+
   /**
    * Cleanup kubernetes engine info by identifier.
    */
@@ -117,4 +121,8 @@ trait MetadataStore extends Closeable {
    * @param limit the maximum number of kubernetes engine info to be cleaned up.
    */
   def cleanupKubernetesEngineInfoByAge(maxAge: Long, limit: Int): Int
+
+  def cleanupKubernetesEngineInfoByAge(maxAge: Long): Int = {
+    cleanupKubernetesEngineInfoByAge(maxAge, Int.MaxValue)
+  }
 }

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

@@ -414,9 +414,10 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
   override def cleanupMetadataByAge(maxAge: Long, limit: Int): Int = {
     val minEndTime = System.currentTimeMillis() - maxAge
     val query =
-      s"DELETE FROM $METADATA_TABLE WHERE state IN ($terminalStates) AND end_time < ? liMIT ?"
+      s"DELETE FROM $METADATA_TABLE WHERE state IN ($terminalStates) AND end_time < ?" +
+        s" ${dialect.deleteFromLimitClause(limit)}"
     JdbcUtils.withConnection { connection =>
-      withUpdateCount(connection, query, minEndTime, limit) { count =>
+      withUpdateCount(connection, query, minEndTime) { count =>
         info(s"Cleaned up $count records older than $maxAge ms from $METADATA_TABLE limit:$limit.")
         count
       }
@@ -467,9 +468,10 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
 
   override def cleanupKubernetesEngineInfoByAge(maxAge: Long, limit: Int): Int = {
     val minUpdateTime = System.currentTimeMillis() - maxAge
-    val query = s"DELETE FROM $KUBERNETES_ENGINE_INFO_TABLE WHERE update_time < ? LIMIT ?"
+    val query = s"DELETE FROM $KUBERNETES_ENGINE_INFO_TABLE WHERE update_time < ?" +
+      s" ${dialect.deleteFromLimitClause(limit)}"
     JdbcUtils.withConnection { connection =>
-      withUpdateCount(connection, query, minUpdateTime, limit) { count =>
+      withUpdateCount(connection, query, minUpdateTime) { count =>
         info(s"Cleaned up $count records older than $maxAge ms from $KUBERNETES_ENGINE_INFO_TABLE" +
           s" limit $limit.")
         count

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

@@ -19,6 +19,7 @@ package org.apache.kyuubi.server.metadata.jdbc
 
 trait JdbcDatabaseDialect {
   def limitClause(limit: Int, offset: Int): String
+  def deleteFromLimitClause(limit: Int): String
   def insertOrReplace(
       table: String,
       cols: Seq[String],
@@ -31,6 +32,10 @@ class GenericDatabaseDialect extends JdbcDatabaseDialect {
     s"LIMIT $limit OFFSET $offset"
   }
 
+  override def deleteFromLimitClause(limit: Int): String = {
+    "" // Generic dialect does not support LIMIT in DELETE statements
+  }
+
   override def insertOrReplace(
       table: String,
       cols: Seq[String],
@@ -71,6 +76,10 @@ class MySQLDatabaseDialect extends GenericDatabaseDialect {
        |${cols.filterNot(_ == keyCol).map(c => s"$c = new.$c").mkString(",")}
        |""".stripMargin
   }
+
+  override def deleteFromLimitClause(limit: Int): String = {
+    s"LIMIT $limit"
+  }
 }
 class PostgreSQLDatabaseDialect extends GenericDatabaseDialect {
   override def insertOrReplace(

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

@@ -214,10 +214,12 @@ class MetadataManagerSuite extends KyuubiFunSuite {
       metadataManager.getBatches(MetadataFilter(), 0, Int.MaxValue).foreach { batch =>
         // close the batch if not ended
         if (batch.getEndTime == 0) {
-          metadataManager.updateMetadata(Metadata(
-            identifier = batch.getId,
-            state = OperationState.CLOSED.toString,
-            endTime = System.currentTimeMillis()), false)
+          metadataManager.updateMetadata(
+            Metadata(
+              identifier = batch.getId,
+              state = OperationState.CLOSED.toString,
+              endTime = System.currentTimeMillis()),
+            false)
         }
       }
 

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

@@ -45,7 +45,7 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
       batch =>
         jdbcMetadataStore.cleanupMetadataByIdentifier(batch.identifier)
     }
-    jdbcMetadataStore.cleanupKubernetesEngineInfoByAge(0, Int.MaxValue)
+    jdbcMetadataStore.cleanupKubernetesEngineInfoByAge(0)
     jdbcMetadataStore.close()
   }
 
@@ -242,7 +242,7 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
       Int.MaxValue).isEmpty)
 
     eventually(Timeout(3.seconds)) {
-      jdbcMetadataStore.cleanupMetadataByAge(1000, Int.MaxValue)
+      jdbcMetadataStore.cleanupMetadataByAge(1000)
       assert(jdbcMetadataStore.getMetadata(batchId) == null)
     }
   }