[TOREE-541] Reply message should implement status field (#207)

pan3793 · web-flow · commit 6bf8c0658207 · 2024-09-01T18:37:26.000-07:00
diff --git a/client/src/main/scala/org/apache/toree/kernel/protocol/v5/client/socket/StdinClient.scala b/client/src/main/scala/org/apache/toree/kernel/protocol/v5/client/socket/StdinClient.scala
@@ -22,7 +22,7 @@ import org.apache.toree.communication.ZMQMessage
 import org.apache.toree.communication.security.SecurityActorType
 import org.apache.toree.kernel.protocol.v5.client.ActorLoader
 import org.apache.toree.kernel.protocol.v5.{HeaderBuilder, KMBuilder, KernelMessage}
-import org.apache.toree.kernel.protocol.v5.content.{InputReply, InputRequest}
+import org.apache.toree.kernel.protocol.v5.content.{InputReply, InputReplyOk, InputRequest}
 import org.apache.toree.utils.LogLike
 import org.apache.toree.kernel.protocol.v5.client.Utilities._
 import play.api.libs.json.Json
@@ -76,7 +76,7 @@ class StdinClient(
         val inputRequest =
           Json.parse(kernelMessage.contentString).as[InputRequest]
         val value = responseFunc(inputRequest.prompt, inputRequest.password)
-        val inputReply = InputReply(value)
+        val inputReply = InputReplyOk(value)
 
         val newKernelMessage = KMBuilder()
           .withParent(kernelMessage)
diff --git a/kernel/src/main/scala/org/apache/toree/kernel/protocol/v5/handler/CommInfoRequestHandler.scala b/kernel/src/main/scala/org/apache/toree/kernel/protocol/v5/handler/CommInfoRequestHandler.scala
@@ -19,13 +19,13 @@ package org.apache.toree.kernel.protocol.v5.handler
 
 import org.apache.toree.comm.CommStorage
 import org.apache.toree.kernel.protocol.v5._
-import org.apache.toree.kernel.protocol.v5.content.CommInfoReply
+import org.apache.toree.kernel.protocol.v5.content.{CommInfoReply, CommInfoReplyOk}
 import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
 import org.apache.toree.utils.MessageLogSupport
 import play.api.libs.json.Json
 
 import scala.concurrent.ExecutionContext.Implicits.global
-import scala.concurrent.{Future, future}
+import scala.concurrent.Future
 
 /**
  * Receives a CommInfoRequest KernelMessage and returns a CommInfoReply
@@ -63,7 +63,7 @@ class CommInfoRequestHandler(
         commStorage.getTargets().map(buildCommMap(_)).reduce(_ ++ _)
       }
     }
-    val commInfoReply = CommInfoReply(commMap.asInstanceOf[Map[String, Map[String, String]]])
+    val commInfoReply = CommInfoReplyOk(commMap.asInstanceOf[Map[String, Map[String, String]]])
 
     val kernelInfo = SparkKernelInfo
 
diff --git a/kernel/src/main/scala/org/apache/toree/kernel/protocol/v5/handler/KernelInfoRequestHandler.scala b/kernel/src/main/scala/org/apache/toree/kernel/protocol/v5/handler/KernelInfoRequestHandler.scala
@@ -18,7 +18,7 @@
 package org.apache.toree.kernel.protocol.v5.handler
 
 import org.apache.toree.kernel.protocol.v5._
-import org.apache.toree.kernel.protocol.v5.content.KernelInfoReply
+import org.apache.toree.kernel.protocol.v5.content.{KernelInfoReply, KernelInfoReplyOk}
 import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
 import org.apache.toree.utils.LogLike
 
@@ -37,7 +37,7 @@ class KernelInfoRequestHandler(actorLoader: ActorLoader, languageInfo: LanguageI
       logger.debug("Sending kernel info reply message")
 
       val kernelInfo = SparkKernelInfo
-      val kernelInfoReply = KernelInfoReply(
+      val kernelInfoReply = KernelInfoReplyOk(
         kernelInfo.protocolVersion,
         kernelInfo.implementation,
         kernelInfo.implementationVersion,
diff --git a/kernel/src/main/scala/org/apache/toree/kernel/protocol/v5/handler/ShutdownHandler.scala b/kernel/src/main/scala/org/apache/toree/kernel/protocol/v5/handler/ShutdownHandler.scala
@@ -18,7 +18,7 @@
 package org.apache.toree.kernel.protocol.v5.handler
 
 import org.apache.toree.comm.{CommRegistrar, CommStorage, KernelCommWriter}
-import org.apache.toree.kernel.protocol.v5.content.{ShutdownReply, ShutdownRequest, CommOpen}
+import org.apache.toree.kernel.protocol.v5.content.{CommOpen, ShutdownReply, ShutdownReplyOk, ShutdownRequest}
 import org.apache.toree.kernel.protocol.v5.kernel.{ActorLoader, Utilities}
 import org.apache.toree.kernel.protocol.v5._
 import org.apache.toree.security.KernelSecurityManager
@@ -43,7 +43,7 @@ class ShutdownHandler(
 
     val kernelInfo = SparkKernelInfo
 
-    val shutdownReply = ShutdownReply(false)
+    val shutdownReply = ShutdownReplyOk(false)
 
     val replyHeader = Header(
       java.util.UUID.randomUUID.toString,
diff --git a/kernel/src/test/scala/org/apache/toree/kernel/protocol/v5/handler/InputRequestReplyHandlerSpec.scala b/kernel/src/test/scala/org/apache/toree/kernel/protocol/v5/handler/InputRequestReplyHandlerSpec.scala
@@ -22,8 +22,8 @@ import java.util.concurrent.ConcurrentHashMap
 import org.apache.pekko.actor.{Props, ActorRef, ActorSystem}
 import org.apache.pekko.testkit.{TestProbe, ImplicitSender, TestKit}
 import org.apache.toree.Main
-import org.apache.toree.kernel.protocol.v5.content.InputReply
-import org.apache.toree.kernel.protocol.v5.{HeaderBuilder, MessageType, KMBuilder, SystemActorType}
+import org.apache.toree.kernel.protocol.v5.content.{InputReply, InputReplyOk}
+import org.apache.toree.kernel.protocol.v5.{HeaderBuilder, KMBuilder, MessageType, SystemActorType}
 import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
 import org.scalatest.concurrent.Eventually
 import org.scalatestplus.mockito.MockitoSugar
@@ -105,7 +105,7 @@ class InputRequestReplyHandlerSpec
             msg_type = MessageType.Incoming.InputReply.toString,
             session = session
           ))
-          .withContentString(InputReply(expected))
+          .withContentString(InputReplyOk(expected))
           .build
 
         // Add our fake sender actor to the receiving end of the message
@@ -127,7 +127,7 @@ class InputRequestReplyHandlerSpec
           msg_type = MessageType.Incoming.InputReply.toString,
           session = session
         ))
-          .withContentString(InputReply(expected))
+          .withContentString(InputReplyOk(expected))
           .build
 
         fakeSender.send(inputRequestReplyHandler, inputReplyMessage)
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/CommInfoReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/CommInfoReply.scala
@@ -21,8 +21,9 @@ import org.apache.toree.kernel.protocol.v5.KernelMessageContent
 import play.api.libs.json.Json
 
 case class CommInfoReply(
+  status: String,
   comms: Map[String, Map[String, String]]
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(CommInfoReply.commInfoReplyWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/CompleteReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/CompleteReply.scala
@@ -29,7 +29,7 @@ case class CompleteReply (
   ename: Option[String],
   evalue: Option[String],
   traceback: Option[List[String]]
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(CompleteReply.completeReplyWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ConnectReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ConnectReply.scala
@@ -20,6 +20,8 @@ package org.apache.toree.kernel.protocol.v5.content
 import org.apache.toree.kernel.protocol.v5.KernelMessageContent
 import play.api.libs.json.Json
 
+// Deprecated since version 5.1: connect_request/reply have not proved useful, and are
+// considered deprecated. Kernels are not expected to implement handlers for this message.
 case class ConnectReply(
   shell_port: Int,
   iopub_port: Int,
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ConnectRequest.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ConnectRequest.scala
@@ -20,6 +20,8 @@ package org.apache.toree.kernel.protocol.v5.content
 import org.apache.toree.kernel.protocol.v5.KernelMessageContent
 import play.api.libs.json._
 
+// Deprecated since version 5.1: connect_request/reply have not proved useful, and are
+// considered deprecated. Kernels are not expected to implement handlers for this message.
 case class ConnectRequest() extends KernelMessageContent {
   override def content : String =
     Json.toJson(this)(ConnectRequest.connectRequestWrites).toString
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ExecuteReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ExecuteReply.scala
@@ -33,7 +33,7 @@ case class ExecuteReply(
   ename: Option[String],
   evalue: Option[String],
   traceback: Option[List[String]]
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
 
   override def content : String =
     Json.toJson(this)(ExecuteReply.executeReplyWrites).toString
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/HistoryReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/HistoryReply.scala
@@ -22,12 +22,13 @@ import play.api.libs.json.Json
 
 
 case class HistoryReply(
+  status: String,
   // TODO: This is really (String, String, String | (String, String)), look
   // TODO: into writing implicits to handle tuples
 
   // NOTE: Currently, only handle (String, String, String)
   history: List[String]
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(HistoryReply.historyReplyWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/InputReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/InputReply.scala
@@ -21,8 +21,9 @@ import org.apache.toree.kernel.protocol.v5.KernelMessageContent
 import play.api.libs.json._
 
 case class InputReply(
+  status: String,
   value: String
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(InputReply.inputReplyWrites).toString()
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/InspectReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/InspectReply.scala
@@ -30,7 +30,7 @@ case class InspectReply(
   ename: Option[String],
   evalue: Option[String],
   traceback: Option[List[String]]
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(InspectReply.inspectReplyOkWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/IsCompleteReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/IsCompleteReply.scala
@@ -17,13 +17,13 @@
 
 package org.apache.toree.kernel.protocol.v5.content
 
-import org.apache.toree.kernel.protocol.v5.{KernelMessageContent}
+import org.apache.toree.kernel.protocol.v5.KernelMessageContent
 import play.api.libs.json.{Json, Reads, Writes}
 
-case class IsCompleteReply (
-                           status: String,
-                           indent: String
-                         ) extends KernelMessageContent {
+case class IsCompleteReply(
+  status: String,
+  indent: String
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(IsCompleteReply.isCompleteReplyWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/KernelInfoReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/KernelInfoReply.scala
@@ -22,12 +22,13 @@ import org.apache.toree.kernel.protocol.v5.LanguageInfo
 import play.api.libs.json.Json
 
 case class KernelInfoReply (
+  status: String,
   protocol_version: String,
   implementation: String,
   implementation_version: String,
   language_info: LanguageInfo,
   banner: String
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content: String =
     Json.toJson(this)(KernelInfoReply.kernelInfoReplyWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ReplyContent.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ReplyContent.scala
@@ -0,0 +1,30 @@
+/*
+ *  Licensed to the Apache Software Foundation (ASF) under one or more
+ *  contributor license agreements.  See the NOTICE file distributed with
+ *  this work for additional information regarding copyright ownership.
+ *  The ASF licenses this file to You under the Apache License, Version 2.0
+ *  (the "License"); you may not use this file except in compliance with
+ *  the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License
+ */
+
+package org.apache.toree.kernel.protocol.v5.content
+
+import org.apache.toree.kernel.protocol.v5.KernelMessageContent
+
+/**
+ * Represents a series of subclasses of KernelMessageContent that embodies the
+ * xyzReply content types.
+ */
+trait ReplyContent {
+  this: KernelMessageContent =>
+
+  def status: String
+}
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ShutdownReply.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/ShutdownReply.scala
@@ -21,8 +21,9 @@ import org.apache.toree.kernel.protocol.v5.KernelMessageContent
 import play.api.libs.json.Json
 
 case class ShutdownReply(
+  status: String,
   restart: Boolean
-) extends KernelMessageContent {
+) extends KernelMessageContent with ReplyContent {
   override def content : String =
     Json.toJson(this)(ShutdownReply.shutdownReplyWrites).toString
 }
diff --git a/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/package.scala b/protocol/src/main/scala/org/apache/toree/kernel/protocol/v5/content/package.scala
@@ -74,4 +74,30 @@ package object content {
     _: List[String], _: Int, _: Int, _: Metadata, "error", _: Option[String],
     _: Option[String], _: Option[List[String]]
   )
+
+  // Provide an InputReplyOk type and object representing a
+  // partially-completed InputReply
+  type InputReplyOk = InputReply
+  val InputReplyOk = InputReply("ok", _: String)
+
+  // Provide an CommInfoReplyOk type and object representing a
+  // partially-completed CommInfoReply
+  type CommInfoReplyOk = CommInfoReply
+  val CommInfoReplyOk = CommInfoReply("ok", _: Map[String, Map[String, String]])
+
+  // Provide an KernelInfoReplyOk type and object representing a
+  // partially-completed KernelInfoReply
+  type KernelInfoReplyOk = KernelInfoReply
+  val KernelInfoReplyOk = KernelInfoReply(
+    "ok", _: String, _: String, _: String, _: LanguageInfo, _: String)
+
+  // Provide an ShutdownReplyOk type and object representing a
+  // partially-completed ShutdownReply
+  type ShutdownReplyOk = ShutdownReply
+  val ShutdownReplyOk = ShutdownReply("ok", _: Boolean)
+
+  // Provide an HistoryReplyOk type and object representing a
+  // partially-completed HistoryReply
+  type HistoryReplyOk = HistoryReply
+  val HistoryReplyOk = HistoryReply("ok", _: List[String])
 }
diff --git a/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/HistoryReplySpec.scala b/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/HistoryReplySpec.scala
@@ -27,11 +27,12 @@ import org.apache.toree.kernel.protocol.v5._
 class HistoryReplySpec extends AnyFunSpec with Matchers {
   val historyReplyJson: JsValue = Json.parse("""
   {
+    "status": "ok",
     "history": ["<STRING>", "<STRING2>", "<STRING3>"]
   }
   """)
 
-  val historyReply = HistoryReply(
+  val historyReply = HistoryReplyOk(
     List("<STRING>", "<STRING2>", "<STRING3>")
   )
 
diff --git a/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/InputReplySpec.scala b/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/InputReplySpec.scala
@@ -25,11 +25,12 @@ import play.api.libs.json._
 class InputReplySpec extends AnyFunSpec with Matchers {
   val inputReplyJson: JsValue = Json.parse("""
   {
+    "status": "ok",
     "value": "<STRING>"
   }
   """)
 
-  val inputReply = InputReply(
+  val inputReply = InputReplyOk(
     "<STRING>"
   )
 
diff --git a/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/KernelInfoReplySpec.scala b/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/KernelInfoReplySpec.scala
@@ -26,6 +26,7 @@ import play.api.libs.json._
 class KernelInfoReplySpec extends AnyFunSpec with Matchers {
   val kernelInfoReplyJson: JsValue = Json.parse("""
   {
+    "status": "ok",
     "protocol_version": "x.y.z",
     "implementation": "<name>",
     "implementation_version": "z.y.x",
@@ -34,7 +35,7 @@ class KernelInfoReplySpec extends AnyFunSpec with Matchers {
   }
   """)
 
-  val kernelInfoReply: KernelInfoReply = KernelInfoReply(
+  val kernelInfoReply: KernelInfoReply = KernelInfoReplyOk(
     "x.y.z", "<name>", "z.y.x", LanguageInfo("<some language>", "a.b.c", Some("<some extension>")), "<some banner>"
   )
 
diff --git a/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/ShutdownReplySpec.scala b/protocol/src/test/scala/org/apache/toree/kernel/protocol/v5/content/ShutdownReplySpec.scala
@@ -25,11 +25,12 @@ import play.api.libs.json._
 class ShutdownReplySpec extends AnyFunSpec with Matchers {
   val shutdownReplyJson: JsValue = Json.parse("""
   {
+    "status": "ok",
     "restart": true
   }
   """)
 
-  val shutdownReply: ShutdownReply = ShutdownReply(
+  val shutdownReply: ShutdownReply = ShutdownReplyOk(
     true
   )
 

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ case class CompleteReply (`
`29`	`29`	`ename: Option[String],`
`30`	`30`	`evalue: Option[String],`
`31`	`31`	`traceback: Option[List[String]]`
`32`		`-) extends KernelMessageContent {`
	`32`	`+) extends KernelMessageContent with ReplyContent {`
`33`	`33`	`override def content : String =`
`34`	`34`	`Json.toJson(this)(CompleteReply.completeReplyWrites).toString`
`35`	`35`	`}`