From cc2e15cb565b1ce55ab7e2015c51059e2250c474 Mon Sep 17 00:00:00 2001
From: Dongjoon Hyun <dongjoon@apache.org>
Date: Mon, 17 Mar 2025 22:20:01 -0700
Subject: [PATCH] [SPARK-51539] Refactor `SparkConnectClient` to use `analyze`
 helper function

---
 Sources/SparkConnect/SparkConnectClient.swift | 21 +++++++++++--------
 Sources/SparkConnect/TypeAliases.swift        |  1 +
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index b9b6e34..c0c0828 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -64,12 +64,9 @@ public actor SparkConnectClient {
 
       self.sessionID = sessionID
       let service = SparkConnectService.Client(wrapping: client)
-      let version = AnalyzePlanRequest.SparkVersion()
-      var request = AnalyzePlanRequest()
-      request.clientType = clientType
-      request.userContext = userContext
-      request.sessionID = self.sessionID!
-      request.analyze = .sparkVersion(version)
+      let request = analyze(self.sessionID!, {
+        return OneOf_Analyze.sparkVersion(AnalyzePlanRequest.SparkVersion())
+      })
       let response = try await service.analyzePlan(request)
       return response
     }
@@ -243,13 +240,19 @@ public actor SparkConnectClient {
   func getAnalyzePlanRequest(_ sessionID: String, _ plan: Plan) async
     -> AnalyzePlanRequest
   {
+    return analyze(sessionID, {
+      var schema = AnalyzePlanRequest.Schema()
+      schema.plan = plan
+      return OneOf_Analyze.schema(schema)
+    })
+  }
+
+  private func analyze(_ sessionID: String, _ f: () -> OneOf_Analyze) -> AnalyzePlanRequest {
     var request = AnalyzePlanRequest()
     request.clientType = clientType
     request.userContext = userContext
     request.sessionID = self.sessionID!
-    var schema = AnalyzePlanRequest.Schema()
-    schema.plan = plan
-    request.analyze = .schema(schema)
+    request.analyze = f()
     return request
   }
 
diff --git a/Sources/SparkConnect/TypeAliases.swift b/Sources/SparkConnect/TypeAliases.swift
index 8e68723..8154662 100644
--- a/Sources/SparkConnect/TypeAliases.swift
+++ b/Sources/SparkConnect/TypeAliases.swift
@@ -25,6 +25,7 @@ typealias Plan = Spark_Connect_Plan
 typealias Project = Spark_Connect_Project
 typealias KeyValue = Spark_Connect_KeyValue
 typealias Limit = Spark_Connect_Limit
+typealias OneOf_Analyze = AnalyzePlanRequest.OneOf_Analyze
 typealias Range = Spark_Connect_Range
 typealias Relation = Spark_Connect_Relation
 typealias SparkConnectService = Spark_Connect_SparkConnectService