[SPARK-42688][CONNECT] Rename Connect proto Request client_id to sess…

…ion_id ### What changes were proposed in this pull request? Rename Connect proto requests `client_id` to `session_id`. On the one hand when I read `client_id` I was confused on what it is used to, even after reading the proto documentation. On the other hand, client sides already use session_id: https://github.com/apache/spark/blob/9bf174f9722e34f13bfaede5e59f989bf2a511e9/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala#L51 https://github.com/apache/spark/blob/9bf174f9722e34f13bfaede5e59f989bf2a511e9/python/pyspark/sql/connect/client.py#L522 ### Why are the changes needed? Code readability ### Does this PR introduce _any_ user-facing change? NO ### How was this patch tested? Existing UT Closes apache#40309 from amaliujia/update_client_id. Authored-by: Rui Wang <[email protected]> Signed-off-by: Herman van Hovell <[email protected]>
saileshbaidya · Mar 7, 2023 · dfdc4a1 · dfdc4a1
1 parent d4818df
commit dfdc4a1
Show file tree

Hide file tree

Showing 13 changed files with 223 additions and 204 deletions.
diff --git a/...ct/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/...ct/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala
@@ -63,7 +63,7 @@ private[sql] class SparkConnectClient(
       .newBuilder()
       .setPlan(plan)
       .setUserContext(userContext)
-      .setClientId(sessionId)
+      .setSessionId(sessionId)
       .setClientType(userAgent)
       .build()
     stub.executePlan(request)
@@ -78,7 +78,7 @@ private[sql] class SparkConnectClient(
     val request = proto.ConfigRequest
       .newBuilder()
       .setOperation(operation)
-      .setClientId(sessionId)
+      .setSessionId(sessionId)
       .setClientType(userAgent)
       .setUserContext(userContext)
       .build()
@@ -157,7 +157,7 @@ private[sql] class SparkConnectClient(
   private def analyze(builder: proto.AnalyzePlanRequest.Builder): proto.AnalyzePlanResponse = {
     val request = builder
       .setUserContext(userContext)
-      .setClientId(sessionId)
+      .setSessionId(sessionId)
       .setClientType(userAgent)
       .build()
     analyze(request)

diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala
@@ -612,8 +612,8 @@ class ClientE2ETestSuite extends RemoteSparkSession {
   }
 
   test("SparkSession newSession") {
-    val oldId = spark.sql("SELECT 1").analyze.getClientId
-    val newId = spark.newSession().sql("SELECT 1").analyze.getClientId
+    val oldId = spark.sql("SELECT 1").analyze.getSessionId
+    val newId = spark.newSession().sql("SELECT 1").analyze.getSessionId
     assert(oldId != newId)
   }
 

diff --git a/...ient/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala b/...ient/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala
@@ -75,11 +75,11 @@ class SparkConnectClientSuite extends ConnectFunSuite with BeforeAndAfterEach {
     client = clientBuilder(server.getPort)
     val request = AnalyzePlanRequest
       .newBuilder()
-      .setClientId("abc123")
+      .setSessionId("abc123")
       .build()
 
     val response = client.analyze(request)
-    assert(response.getClientId === "abc123")
+    assert(response.getSessionId === "abc123")
   }
 
   test("Test connection") {
@@ -99,7 +99,7 @@ class SparkConnectClientSuite extends ConnectFunSuite with BeforeAndAfterEach {
       .connectionString(s"sc://localhost:${server.getPort}/;use_ssl=true")
       .build()
 
-    val request = AnalyzePlanRequest.newBuilder().setClientId("abc123").build()
+    val request = AnalyzePlanRequest.newBuilder().setSessionId("abc123").build()
 
     // Failed the ssl handshake as the dummy server does not have any server credentials installed.
     assertThrows[StatusRuntimeException] {
@@ -201,11 +201,11 @@ class DummySparkConnectService() extends SparkConnectServiceGrpc.SparkConnectSer
       request: ExecutePlanRequest,
       responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
     // Reply with a dummy response using the same client ID
-    val requestClientId = request.getClientId
+    val requestSessionId = request.getSessionId
     inputPlan = request.getPlan
     val response = ExecutePlanResponse
       .newBuilder()
-      .setClientId(requestClientId)
+      .setSessionId(requestSessionId)
       .build()
     responseObserver.onNext(response)
     responseObserver.onCompleted()
@@ -215,7 +215,7 @@ class DummySparkConnectService() extends SparkConnectServiceGrpc.SparkConnectSer
       request: AnalyzePlanRequest,
       responseObserver: StreamObserver[AnalyzePlanResponse]): Unit = {
     // Reply with a dummy response using the same client ID
-    val requestClientId = request.getClientId
+    val requestSessionId = request.getSessionId
     request.getAnalyzeCase match {
       case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
         inputPlan = request.getSchema.getPlan
@@ -233,7 +233,7 @@ class DummySparkConnectService() extends SparkConnectServiceGrpc.SparkConnectSer
     }
     val response = AnalyzePlanResponse
       .newBuilder()
-      .setClientId(requestClientId)
+      .setSessionId(requestSessionId)
       .build()
     responseObserver.onNext(response)
     responseObserver.onCompleted()

diff --git a/connector/connect/common/src/main/protobuf/spark/connect/base.proto b/connector/connect/common/src/main/protobuf/spark/connect/base.proto
@@ -58,9 +58,10 @@ message UserContext {
 message AnalyzePlanRequest {
   // (Required)
   //
-  // The client_id is set by the client to be able to collate streaming responses from
-  // different queries.
-  string client_id = 1;
+  // The session_id specifies a spark session for a user id (which is specified
+  // by user_context.user_id). The session_id is set by the client to be able to
+  // collate streaming responses from different queries within the dedicated session.
+  string session_id = 1;
 
   // (Required) User context
   UserContext user_context = 2;
@@ -161,7 +162,7 @@ message AnalyzePlanRequest {
 // Response to performing analysis of the query. Contains relevant metadata to be able to
 // reason about the performance.
 message AnalyzePlanResponse {
-  string client_id = 1;
+  string session_id = 1;
 
   oneof result {
     Schema schema = 2;
@@ -217,11 +218,15 @@ message AnalyzePlanResponse {
 message ExecutePlanRequest {
   // (Required)
   //
-  // The client_id is set by the client to be able to collate streaming responses from
-  // different queries.
-  string client_id = 1;
+  // The session_id specifies a spark session for a user id (which is specified
+  // by user_context.user_id). The session_id is set by the client to be able to
+  // collate streaming responses from different queries within the dedicated session.
+  string session_id = 1;
 
   // (Required) User context
+  //
+  // user_context.user_id and session+id both identify a unique remote spark session on the
+  // server side.
   UserContext user_context = 2;
 
   // (Required) The logical plan to be executed / analyzed.
@@ -234,9 +239,9 @@ message ExecutePlanRequest {
 }
 
 // The response of a query, can be one or more for each request. Responses belonging to the
-// same input query, carry the same `client_id`.
+// same input query, carry the same `session_id`.
 message ExecutePlanResponse {
-  string client_id = 1;
+  string session_id = 1;
 
   // Union type for the different response messages.
   oneof response_type {
@@ -304,9 +309,10 @@ message KeyValue {
 message ConfigRequest {
   // (Required)
   //
-  // The client_id is set by the client to be able to collate streaming responses from
-  // different queries.
-  string client_id = 1;
+  // The session_id specifies a spark session for a user id (which is specified
+  // by user_context.user_id). The session_id is set by the client to be able to
+  // collate streaming responses from different queries within the dedicated session.
+  string session_id = 1;
 
   // (Required) User context
   UserContext user_context = 2;
@@ -369,7 +375,7 @@ message ConfigRequest {
 
 // Response to the config request.
 message ConfigResponse {
-  string client_id = 1;
+  string session_id = 1;
 
   // (Optional) The result key-value pairs.
   //
@@ -386,9 +392,12 @@ message ConfigResponse {
 // Request to transfer client-local artifacts.
 message AddArtifactsRequest {
 
-  // The client_id is set by the client to be able to collate streaming responses from
-  // different queries.
-  string client_id = 1;
+  // (Required)
+  //
+  // The session_id specifies a spark session for a user id (which is specified
+  // by user_context.user_id). The session_id is set by the client to be able to
+  // collate streaming responses from different queries within the dedicated session.
+  string session_id = 1;
 
   // User context
   UserContext user_context = 2;

diff --git a/...nect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/...nect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -1459,7 +1459,7 @@ class SparkConnectPlanner(val session: SparkSession) {
 
   def process(
       command: proto.Command,
-      clientId: String,
+      sessionId: String,
       responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
     command.getCommandTypeCase match {
       case proto.Command.CommandTypeCase.REGISTER_FUNCTION =>
@@ -1473,14 +1473,14 @@ class SparkConnectPlanner(val session: SparkSession) {
       case proto.Command.CommandTypeCase.EXTENSION =>
         handleCommandPlugin(command.getExtension)
       case proto.Command.CommandTypeCase.SQL_COMMAND =>
-        handleSqlCommand(command.getSqlCommand, clientId, responseObserver)
+        handleSqlCommand(command.getSqlCommand, sessionId, responseObserver)
       case _ => throw new UnsupportedOperationException(s"$command not supported.")
     }
   }
 
   def handleSqlCommand(
       getSqlCommand: SqlCommand,
-      clientId: String,
+      sessionId: String,
       responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
     // Eagerly execute commands of the provided SQL string.
     val df = session.sql(getSqlCommand.getSql, getSqlCommand.getArgsMap)
@@ -1537,12 +1537,12 @@ class SparkConnectPlanner(val session: SparkSession) {
     responseObserver.onNext(
       ExecutePlanResponse
         .newBuilder()
-        .setClientId(clientId)
+        .setSessionId(sessionId)
         .setSqlCommandResult(result)
         .build())
 
     // Send Metrics
-    SparkConnectStreamHandler.sendMetricsToResponse(clientId, df)
+    SparkConnectStreamHandler.sendMetricsToResponse(sessionId, df)
   }
 
   private def handleRegisterUserDefinedFunction(

diff --git a/...rver/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAnalyzeHandler.scala b/...rver/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAnalyzeHandler.scala
@@ -35,7 +35,7 @@ private[connect] class SparkConnectAnalyzeHandler(
   def handle(request: proto.AnalyzePlanRequest): Unit = {
     val session =
       SparkConnectService
-        .getOrCreateIsolatedSession(request.getUserContext.getUserId, request.getClientId)
+        .getOrCreateIsolatedSession(request.getUserContext.getUserId, request.getSessionId)
         .session
     session.withActive {
       val response = process(request, session)
@@ -155,7 +155,7 @@ private[connect] class SparkConnectAnalyzeHandler(
       case other => throw InvalidPlanInput(s"Unknown Analyze Method $other!")
     }
 
-    builder.setClientId(request.getClientId)
+    builder.setSessionId(request.getSessionId)
     builder.build()
   }
 }
diff --git a/...erver/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala b/...erver/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala
@@ -32,7 +32,7 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes
   def handle(request: proto.ConfigRequest): Unit = {
     val session =
       SparkConnectService
-        .getOrCreateIsolatedSession(request.getUserContext.getUserId, request.getClientId)
+        .getOrCreateIsolatedSession(request.getUserContext.getUserId, request.getSessionId)
         .session
 
     val builder = request.getOperation.getOpTypeCase match {
@@ -53,7 +53,7 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes
       case _ => throw new UnsupportedOperationException(s"${request.getOperation} not supported.")
     }
 
-    builder.setClientId(request.getClientId)
+    builder.setSessionId(request.getSessionId)
     responseObserver.onNext(builder.build())
     responseObserver.onCompleted()
   }

diff --git a/...erver/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala b/...erver/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala
@@ -44,7 +44,7 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[ExecutePlanResp
   def handle(v: ExecutePlanRequest): Unit = {
     val session =
       SparkConnectService
-        .getOrCreateIsolatedSession(v.getUserContext.getUserId, v.getClientId)
+        .getOrCreateIsolatedSession(v.getUserContext.getUserId, v.getSessionId)
         .session
     session.withActive {
       v.getPlan.getOpTypeCase match {
@@ -60,20 +60,20 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[ExecutePlanResp
     // Extract the plan from the request and convert it to a logical plan
     val planner = new SparkConnectPlanner(session)
     val dataframe = Dataset.ofRows(session, planner.transformRelation(request.getPlan.getRoot))
-    processAsArrowBatches(request.getClientId, dataframe, responseObserver)
+    processAsArrowBatches(request.getSessionId, dataframe, responseObserver)
     responseObserver.onNext(
-      SparkConnectStreamHandler.sendMetricsToResponse(request.getClientId, dataframe))
+      SparkConnectStreamHandler.sendMetricsToResponse(request.getSessionId, dataframe))
     if (dataframe.queryExecution.observedMetrics.nonEmpty) {
       responseObserver.onNext(
-        SparkConnectStreamHandler.sendObservedMetricsToResponse(request.getClientId, dataframe))
+        SparkConnectStreamHandler.sendObservedMetricsToResponse(request.getSessionId, dataframe))
     }
     responseObserver.onCompleted()
   }
 
   private def handleCommand(session: SparkSession, request: ExecutePlanRequest): Unit = {
     val command = request.getPlan.getCommand
     val planner = new SparkConnectPlanner(session)
-    planner.process(command, request.getClientId, responseObserver)
+    planner.process(command, request.getSessionId, responseObserver)
     responseObserver.onCompleted()
   }
 }
@@ -96,7 +96,7 @@ object SparkConnectStreamHandler {
   }
 
   def processAsArrowBatches(
-      clientId: String,
+      sessionId: String,
       dataframe: DataFrame,
       responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
     val spark = dataframe.sparkSession
@@ -173,7 +173,7 @@ object SparkConnectStreamHandler {
           }
 
           partition.foreach { case (bytes, count) =>
-            val response = proto.ExecutePlanResponse.newBuilder().setClientId(clientId)
+            val response = proto.ExecutePlanResponse.newBuilder().setSessionId(sessionId)
             val batch = proto.ExecutePlanResponse.ArrowBatch
               .newBuilder()
               .setRowCount(count)
@@ -191,7 +191,7 @@ object SparkConnectStreamHandler {
       // Make sure at least 1 batch will be sent.
       if (numSent == 0) {
         val bytes = ArrowConverters.createEmptyArrowBatch(schema, timeZoneId)
-        val response = proto.ExecutePlanResponse.newBuilder().setClientId(clientId)
+        val response = proto.ExecutePlanResponse.newBuilder().setSessionId(sessionId)
         val batch = proto.ExecutePlanResponse.ArrowBatch
           .newBuilder()
           .setRowCount(0L)
@@ -203,17 +203,17 @@ object SparkConnectStreamHandler {
     }
   }
 
-  def sendMetricsToResponse(clientId: String, rows: DataFrame): ExecutePlanResponse = {
+  def sendMetricsToResponse(sessionId: String, rows: DataFrame): ExecutePlanResponse = {
     // Send a last batch with the metrics
     ExecutePlanResponse
       .newBuilder()
-      .setClientId(clientId)
+      .setSessionId(sessionId)
       .setMetrics(MetricGenerator.buildMetrics(rows.queryExecution.executedPlan))
       .build()
   }
 
   def sendObservedMetricsToResponse(
-      clientId: String,
+      sessionId: String,
       dataframe: DataFrame): ExecutePlanResponse = {
     val observedMetrics = dataframe.queryExecution.observedMetrics.map { case (name, row) =>
       val cols = (0 until row.length).map(i => toConnectProtoValue(row(i)))
@@ -226,7 +226,7 @@ object SparkConnectStreamHandler {
     // Prepare a response with the observed metrics.
     ExecutePlanResponse
       .newBuilder()
-      .setClientId(clientId)
+      .setSessionId(sessionId)
       .addAllObservedMetrics(observedMetrics.asJava)
       .build()
   }

diff --git a/...server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala b/...server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala
@@ -221,7 +221,7 @@ class SparkConnectServiceSuite extends SharedSparkSession {
       .newBuilder()
       .setPlan(plan)
       .setUserContext(context)
-      .setClientId("session")
+      .setSessionId("session")
       .build()
 
     // The observer is executed inside this thread. So