handle session exception for all kyuubi operations

save Revert "save" This reverts commit c3f8fca. session
apache · Nov 18, 2023 · 5d9d4dc · 5d9d4dc
1 parent 3478fc9
commit 5d9d4dc
Show file tree

Hide file tree

Showing 14 changed files with 21 additions and 14 deletions.
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala
@@ -182,7 +182,7 @@ class BatchJobSubmission(
     OperationLog.removeCurrentOperationLog()
   }
 
-  override protected def runInternal(): Unit = session.handleSessionException {
+  override protected def runKyuubiOperationInternal(): Unit = {
     val asyncOperation: Runnable = () => {
       try {
         metadata match {

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala
@@ -164,7 +164,7 @@ class ExecuteStatement(
     }
   }
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     if (isTimeoutMonitorEnabled) {
       addTimeoutMonitor(queryTimeout)
     }

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala
@@ -49,7 +49,7 @@ class ExecutedCommandExec(
     OperationLog.removeCurrentOperationLog()
   }
 
-  override protected def runInternal(): Unit = session.handleSessionException {
+  override protected def runKyuubiOperationInternal(): Unit = {
     val asyncOperation: Runnable = () => {
       setState(OperationState.RUNNING)
       try {

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCatalogs.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCatalogs.scala
@@ -21,7 +21,7 @@ import org.apache.kyuubi.session.Session
 
 class GetCatalogs(session: Session) extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getCatalogs
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetColumns.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetColumns.scala
@@ -27,7 +27,7 @@ class GetColumns(
     columnName: String)
   extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getColumns(catalogName, schemaName, tableName, columnName)
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCrossReference.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetCrossReference.scala
@@ -29,7 +29,7 @@ class GetCrossReference(
     foreignTable: String)
   extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getCrossReference(
         primaryCatalog,

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetFunctions.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetFunctions.scala
@@ -26,7 +26,7 @@ class GetFunctions(
     functionName: String)
   extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getFunctions(catalogName, schemaName, functionName)
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetPrimaryKeys.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetPrimaryKeys.scala
@@ -26,7 +26,7 @@ class GetPrimaryKeys(
     tableName: String)
   extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getPrimaryKeys(catalogName, schemaName, tableName)
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetSchemas.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetSchemas.scala
@@ -25,7 +25,7 @@ class GetSchemas(
     schemaName: String)
   extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getSchemas(catalogName, schemaName)
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetTableTypes.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetTableTypes.scala
@@ -21,7 +21,7 @@ import org.apache.kyuubi.session.Session
 
 class GetTableTypes(session: Session) extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getTableTypes
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetTables.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetTables.scala
@@ -27,7 +27,7 @@ class GetTables(
     tableTypes: java.util.List[String])
   extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getTables(catalogName, schemaName, tableName, tableTypes)
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetTypeInfo.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/GetTypeInfo.scala
@@ -21,7 +21,7 @@ import org.apache.kyuubi.session.Session
 
 class GetTypeInfo(session: Session) extends KyuubiOperation(session) {
 
-  override protected def runInternal(): Unit = {
+  override protected def runKyuubiOperationInternal(): Unit = {
     try {
       _remoteOpHandle = client.getTypeInfo
     } catch onError()

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
@@ -32,7 +32,7 @@ import org.apache.kyuubi.metrics.MetricsConstants.{OPERATION_FAIL, OPERATION_OPE
 import org.apache.kyuubi.metrics.MetricsSystem
 import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
 import org.apache.kyuubi.operation.OperationState.OperationState
-import org.apache.kyuubi.session.{KyuubiSessionImpl, KyuubiSessionManager, Session}
+import org.apache.kyuubi.session.{KyuubiSession, KyuubiSessionImpl, KyuubiSessionManager, Session}
 import org.apache.kyuubi.util.ThriftUtils
 
 abstract class KyuubiOperation(session: Session) extends AbstractOperation(session) {
@@ -100,6 +100,13 @@ abstract class KyuubiOperation(session: Session) extends AbstractOperation(sessi
       }
   }
 
+  final override protected def runInternal(): Unit =
+    session.asInstanceOf[KyuubiSession].handleSessionException {
+      runKyuubiOperationInternal()
+    }
+
+  protected def runKyuubiOperationInternal(): Unit
+
   override protected def beforeRun(): Unit = {
     setHasResultSet(true)
     setState(OperationState.RUNNING)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/LaunchEngine.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/LaunchEngine.scala
@@ -53,7 +53,7 @@ class LaunchEngine(session: KyuubiSessionImpl, override val shouldRunAsync: Bool
     OperationLog.removeCurrentOperationLog()
   }
 
-  override protected def runInternal(): Unit = session.handleSessionException {
+  override protected def runKyuubiOperationInternal(): Unit = {
     val asyncOperation: Runnable = () => {
       setState(OperationState.RUNNING)
       try {