add batch_request_args to rest object batch

apache · Dec 7, 2023 · 7438751 · 7438751
1 parent e779b42
commit 7438751
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 5 deletions.
diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/Batch.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/Batch.java
@@ -18,6 +18,7 @@
 package org.apache.kyuubi.client.api.v1.dto;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import org.apache.commons.lang3.builder.ReflectionToStringBuilder;
@@ -38,6 +39,7 @@ public class Batch {
   private long createTime;
   private long endTime;
   private Map<String, String> batchInfo = Collections.emptyMap();
+  private List<String> args = Collections.emptyList();
 
   public Batch() {}
 
@@ -55,7 +57,8 @@ public Batch(
       String state,
       long createTime,
       long endTime,
-      Map<String, String> batchInfo) {
+      Map<String, String> batchInfo,
+      List<String> args) {
     this.id = id;
     this.user = user;
     this.batchType = batchType;
@@ -70,6 +73,7 @@ public Batch(
     this.createTime = createTime;
     this.endTime = endTime;
     this.batchInfo = batchInfo;
+    this.args = args;
   }
 
   public String getId() {
@@ -183,6 +187,17 @@ public Map<String, String> getBatchInfo() {
     return batchInfo;
   }
 
+  public List<String> getArgs() {
+    if (args == null) {
+      return Collections.emptyList();
+    }
+    return args;
+  }
+
+  public void setArgs(List<String> args) {
+    this.args = args;
+  }
+
   public void setBatchInfo(Map<String, String> batchInfo) {
     this.batchInfo = batchInfo;
   }

diff --git a/kyuubi-rest-client/src/test/java/org/apache/kyuubi/client/RestClientTestUtils.java b/kyuubi-rest-client/src/test/java/org/apache/kyuubi/client/RestClientTestUtils.java
@@ -59,7 +59,8 @@ public static Batch generateTestBatch(String id) {
         "RUNNING",
         BATCH_CREATE_TIME,
         0,
-        Collections.emptyMap());
+        Collections.emptyMap(),
+        Collections.emptyList());
   }
 
   public static BatchRequest generateTestBatchRequest() {

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
@@ -109,7 +109,8 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
       batchOpStatus.state.toString,
       session.createTime,
       batchOpStatus.completed,
-      Map.empty[String, String].asJava)
+      Map.empty[String, String].asJava,
+      session.getBatchArgs.asJava)
   }
 
   private def buildBatch(
@@ -147,7 +148,8 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
         currentBatchState,
         metadata.createTime,
         metadata.endTime,
-        Map.empty[String, String].asJava)
+        Map.empty[String, String].asJava,
+        metadata.requestArgs.asJava)
     }.getOrElse(MetadataManager.buildBatch(metadata))
   }
 

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala
@@ -345,6 +345,7 @@ object MetadataManager extends Logging {
       batchState,
       batchMetadata.createTime,
       batchMetadata.endTime,
-      Map.empty[String, String].asJava)
+      Map.empty[String, String].asJava,
+      batchMetadata.requestArgs.asJava)
   }
 }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiBatchSession.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiBatchSession.scala
@@ -61,6 +61,8 @@ class KyuubiBatchSession(
 
   override def createTime: Long = metadata.map(_.createTime).getOrElse(super.createTime)
 
+  def getBatchArgs: Seq[String] = batchArgs
+
   override def getNoOperationTime: Long = {
     if (batchJobSubmissionOp != null) {
       val batchStatus = batchJobSubmissionOp.getStatus