Refine

Author: turboFei

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

@@ -526,8 +526,7 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
   override def upsertKubernetesEngineInfo(engineInfo: KubernetesEngineInfo): Unit = {
     dialect.insertOrReplace(
       KUBERNETES_ENGINE_INFO_TABLE,
-      KUBERNETES_ENGINE_INFO_COLUMNS_TO_INSERT,
-      KUBERNETES_ENGINE_INFO_COLUMNS_TO_REPLACE,
+      KUBERNETES_ENGINE_INFO_COLUMNS,
       KUBERNETES_ENGINE_INFO_KEY_COLUMN) match {
       case Some(query) =>
         JdbcUtils.withConnection { connection =>
@@ -554,7 +553,7 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
 
   override def getKubernetesMetaEngineInfo(identifier: String): KubernetesEngineInfo = {
     val query =
-      s"SELECT $KUBERNETES_ENGINE_INFO_COLUMNS FROM" +
+      s"SELECT $KUBERNETES_ENGINE_INFO_COLUMNS_STR FROM" +
         s" $KUBERNETES_ENGINE_INFO_TABLE WHERE identifier = ?"
     JdbcUtils.withConnection { connection =>
       withResultSet(connection, query, identifier) { rs =>
@@ -807,19 +806,8 @@ object JDBCMetadataStore {
     "peer_instance_closed").mkString(",")
   private val KUBERNETES_ENGINE_INFO_TABLE = "k8s_engine_info"
   private val KUBERNETES_ENGINE_INFO_KEY_COLUMN = "identifier"
-  private val KUBERNETES_ENGINE_INFO_COLUMNS_TO_INSERT = Seq(
-    "identifier",
-    "context",
-    "namespace",
-    "pod_name",
-    "pod_state",
-    "container_state",
-    "engine_id",
-    "engine_name",
-    "engine_state",
-    "engine_error",
-    "update_time")
-  private val KUBERNETES_ENGINE_INFO_COLUMNS_TO_REPLACE = Seq(
+  private val KUBERNETES_ENGINE_INFO_COLUMNS = Seq(
+    KUBERNETES_ENGINE_INFO_KEY_COLUMN,
     "context",
     "namespace",
     "pod_name",
@@ -830,7 +818,5 @@ object JDBCMetadataStore {
     "engine_state",
     "engine_error",
     "update_time")
-  private val KUBERNETES_ENGINE_INFO_COLUMNS =
-    KUBERNETES_ENGINE_INFO_COLUMNS_TO_INSERT.mkString(",")
-
+  private val KUBERNETES_ENGINE_INFO_COLUMNS_STR = KUBERNETES_ENGINE_INFO_COLUMNS.mkString(",")
 }

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

@@ -21,8 +21,7 @@ trait JdbcDatabaseDialect {
   def limitClause(limit: Int, offset: Int): String
   def insertOrReplace(
       table: String,
-      colsToInsert: Seq[String],
-      colsToReplace: Seq[String],
+      cols: Seq[String],
       keyCol: String): Option[String] = None
 }
 
@@ -35,44 +34,41 @@ class GenericDatabaseDialect extends JdbcDatabaseDialect {
 class SQLiteDatabaseDialect extends GenericDatabaseDialect {
   override def insertOrReplace(
       table: String,
-      colsToInsert: Seq[String],
-      colsToReplace: Seq[String],
+      cols: Seq[String],
       keyCol: String): Option[String] = {
     Some(
       s"""
-         |INSERT OR REPLACE INTO $table (${colsToInsert.mkString(",")})
-         |VALUES (${colsToInsert.map(_ => "?").mkString(",")})
+         |INSERT OR REPLACE INTO $table (${cols.mkString(",")})
+         |VALUES (${cols.map(_ => "?").mkString(",")})
          |""".stripMargin)
   }
 }
 class MySQLDatabaseDialect extends GenericDatabaseDialect {
   override def insertOrReplace(
       table: String,
-      colsToInsert: Seq[String],
-      colsToReplace: Seq[String],
+      cols: Seq[String],
       keyCol: String): Option[String] = {
     Some(
       s"""
-         |INSERT INTO $table (${colsToInsert.mkString(",")})
-         |VALUES (${colsToInsert.map(_ => "?").mkString(",")}) AS new
+         |INSERT INTO $table (${cols.mkString(",")})
+         |VALUES (${cols.map(_ => "?").mkString(",")}) AS new
          |ON DUPLICATE KEY UPDATE
-         |${colsToReplace.map(c => s"$c = new.$c").mkString(",")}
+         |${cols.filterNot(_ == keyCol).map(c => s"$c = new.$c").mkString(",")}
          |""".stripMargin)
   }
 }
 class PostgreSQLDatabaseDialect extends GenericDatabaseDialect {
   override def insertOrReplace(
       table: String,
-      colsToInsert: Seq[String],
-      colsToReplace: Seq[String],
+      cols: Seq[String],
       keyCol: String): Option[String] = {
     Some(
       s"""
-         |INSERT INTO $table (${colsToInsert.mkString(",")})
-         |VALUES (${colsToInsert.map(_ => "?").mkString(",")})
+         |INSERT INTO $table (${cols.mkString(",")})
+         |VALUES (${cols.map(_ => "?").mkString(",")})
          |ON CONFLICT ($keyCol)
          |DO UPDATE SET
-         |${colsToReplace.map(c => s"$c = EXCLUDED.$c").mkString(",")}
+         |${cols.filterNot(_ == keyCol).map(c => s"$c = EXCLUDED.$c").mkString(",")}
          |""".stripMargin)
   }
 }