From 4dc6e1df7a935e0d5b799da12a9544dfe0011ed4 Mon Sep 17 00:00:00 2001
From: fwang12 <fwang12@ebay.com>
Date: Wed, 15 Nov 2023 11:19:02 +0800
Subject: [PATCH] [REST] Return remote session id and remote operation id

---
 .../kyuubi/client/api/v1/dto/OperationData.java       | 11 +++++++++++
 .../apache/kyuubi/client/api/v1/dto/SessionData.java  | 11 +++++++++++
 .../scala/org/apache/kyuubi/server/api/ApiUtils.scala |  2 ++
 3 files changed, 24 insertions(+)

diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/OperationData.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/OperationData.java
index 70c2dd3f3a1..dd3b302640b 100644
--- a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/OperationData.java
+++ b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/OperationData.java
@@ -25,6 +25,7 @@
 
 public class OperationData {
   private String identifier;
+  private String remoteId;
   private String statement;
   private String state;
   private Long createTime;
@@ -41,6 +42,7 @@ public OperationData() {}
 
   public OperationData(
       String identifier,
+      String remoteId,
       String statement,
       String state,
       Long createTime,
@@ -53,6 +55,7 @@ public OperationData(
       String kyuubiInstance,
       Map<String, String> metrics) {
     this.identifier = identifier;
+    this.remoteId = remoteId;
     this.statement = statement;
     this.state = state;
     this.createTime = createTime;
@@ -74,6 +77,14 @@ public void setIdentifier(String identifier) {
     this.identifier = identifier;
   }
 
+  public String getRemoteId() {
+    return remoteId;
+  }
+
+  public void setRemoteId(String remoteId) {
+    this.remoteId = remoteId;
+  }
+
   public String getStatement() {
     return statement;
   }
diff --git a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/SessionData.java b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/SessionData.java
index ae7dfdec984..30a4eb51540 100644
--- a/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/SessionData.java
+++ b/kyuubi-rest-client/src/main/java/org/apache/kyuubi/client/api/v1/dto/SessionData.java
@@ -25,6 +25,7 @@
 
 public class SessionData {
   private String identifier;
+  private String remoteId;
   private String user;
   private String ipAddr;
   private Map<String, String> conf;
@@ -40,6 +41,7 @@ public SessionData() {}
 
   public SessionData(
       String identifier,
+      String remoteId,
       String user,
       String ipAddr,
       Map<String, String> conf,
@@ -51,6 +53,7 @@ public SessionData(
       String kyuubiInstance,
       String engineId) {
     this.identifier = identifier;
+    this.remoteId = remoteId;
     this.user = user;
     this.ipAddr = ipAddr;
     this.conf = conf;
@@ -71,6 +74,14 @@ public void setIdentifier(String identifier) {
     this.identifier = identifier;
   }
 
+  public String getRemoteId() {
+    return remoteId;
+  }
+
+  public void setRemoteId(String remoteId) {
+    this.remoteId = remoteId;
+  }
+
   public String getUser() {
     return user;
   }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/ApiUtils.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/ApiUtils.scala
index 5aaf4d7780f..04a12623575 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/ApiUtils.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/ApiUtils.scala
@@ -32,6 +32,7 @@ object ApiUtils extends Logging {
     val sessionEvent = session.getSessionEvent
     new SessionData(
       session.handle.identifier.toString,
+      sessionEvent.map(_.remoteSessionId).getOrElse(""),
       session.user,
       session.ipAddress,
       session.conf.asJava,
@@ -48,6 +49,7 @@ object ApiUtils extends Logging {
     val opEvent = KyuubiOperationEvent(operation)
     new OperationData(
       opEvent.statementId,
+      opEvent.remoteId,
       opEvent.statement,
       opEvent.state,
       opEvent.createTime,