From e53da12419ec34dd134e80f2c96c9619e48cda7e Mon Sep 17 00:00:00 2001
From: fwang12 <fwang12@ebay.com>
Date: Wed, 15 Mar 2023 11:13:07 +0800
Subject: [PATCH] [KYUUBI #4516] Return kyuubi instance in SessionData and
 OperationData

### _Why are the changes needed?_

It is helpful, because now kyuubi gateway does not support multiple nodes ha.

### _How was this patch tested?_
- [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [x] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #4516 from turboFei/kyuubi_instance.

Closes #4516

a679dc524 [fwang12] return kyuubi instance

Authored-by: fwang12 <fwang12@ebay.com>
Signed-off-by: fwang12 <fwang12@ebay.com>
---
 .../kyuubi/client/api/v1/dto/OperationData.java     | 13 ++++++++++++-
 .../kyuubi/client/api/v1/dto/SessionData.java       | 13 ++++++++++++-
 .../org/apache/kyuubi/server/api/ApiUtils.scala     |  6 ++++--
 3 files changed, 28 insertions(+), 4 deletions(-)

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 bc7af4ac9..1b99bb2c6 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
@@ -32,6 +32,7 @@ public class OperationData {
   private String sessionId;
   private String sessionUser;
   private String sessionType;
+  private String kyuubiInstance;
 
   public OperationData() {}
 
@@ -45,7 +46,8 @@ public class OperationData {
       String exception,
       String sessionId,
       String sessionUser,
-      String sessionType) {
+      String sessionType,
+      String kyuubiInstance) {
     this.identifier = identifier;
     this.statement = statement;
     this.state = state;
@@ -56,6 +58,7 @@ public class OperationData {
     this.sessionId = sessionId;
     this.sessionUser = sessionUser;
     this.sessionType = sessionType;
+    this.kyuubiInstance = kyuubiInstance;
   }
 
   public String getIdentifier() {
@@ -138,6 +141,14 @@ public class OperationData {
     this.sessionType = sessionType;
   }
 
+  public String getKyuubiInstance() {
+    return kyuubiInstance;
+  }
+
+  public void setKyuubiInstance(String kyuubiInstance) {
+    this.kyuubiInstance = kyuubiInstance;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;
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 97da4d942..ce5eaae50 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
@@ -33,6 +33,7 @@ public class SessionData {
   private Long idleTime;
   private String exception;
   private String sessionType;
+  private String kyuubiInstance;
 
   public SessionData() {}
 
@@ -45,7 +46,8 @@ public class SessionData {
       Long duration,
       Long idleTime,
       String exception,
-      String sessionType) {
+      String sessionType,
+      String kyuubiInstance) {
     this.identifier = identifier;
     this.user = user;
     this.ipAddr = ipAddr;
@@ -55,6 +57,7 @@ public class SessionData {
     this.idleTime = idleTime;
     this.exception = exception;
     this.sessionType = sessionType;
+    this.kyuubiInstance = kyuubiInstance;
   }
 
   public String getIdentifier() {
@@ -132,6 +135,14 @@ public class SessionData {
     this.sessionType = sessionType;
   }
 
+  public String getKyuubiInstance() {
+    return kyuubiInstance;
+  }
+
+  public void setKyuubiInstance(String kyuubiInstance) {
+    this.kyuubiInstance = kyuubiInstance;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;
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 7a0582613..3e7e36eac 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
@@ -37,7 +37,8 @@ object ApiUtils {
       session.lastAccessTime - session.createTime,
       session.getNoOperationTime,
       session.getSessionEvent.flatMap(_.exception).map(Utils.prettyPrint).getOrElse(""),
-      session.sessionType.toString)
+      session.sessionType.toString,
+      session.connectionUrl)
   }
 
   def operationData(operation: KyuubiOperation): OperationData = {
@@ -52,6 +53,7 @@ object ApiUtils {
       opEvent.exception.map(Utils.prettyPrint).getOrElse(""),
       opEvent.sessionId,
       opEvent.sessionUser,
-      opEvent.sessionType)
+      opEvent.sessionType,
+      operation.getSession.asInstanceOf[KyuubiSession].connectionUrl)
   }
 }