Skip to content

Commit c43bc16

Browse files
committed
[SPARK-52094] Support session_id parameter
### What changes were proposed in this pull request? This PR aims to support `session_id` parameter. ### Why are the changes needed? For feature parity. After this PR, a user can reuse the session id. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass the CIs. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #139 from dongjoon-hyun/SPARK-52094. Authored-by: Dongjoon Hyun <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>
1 parent 2e8324a commit c43bc16

File tree

3 files changed

+23
-1
lines changed

3 files changed

+23
-1
lines changed

Sources/SparkConnect/SparkConnectClient.swift

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@ public actor SparkConnectClient {
5252
for param in self.url.path.split(separator: ";").dropFirst().filter({ !$0.isEmpty }) {
5353
let kv = param.split(separator: "=")
5454
switch String(kv[0]).lowercased() {
55+
case URIParams.PARAM_SESSION_ID:
56+
// SparkSession handles this.
57+
break
5558
case URIParams.PARAM_USER_AGENT:
5659
clientType = String(kv[1])
5760
case URIParams.PARAM_TOKEN:

Sources/SparkConnect/SparkSession.swift

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,19 @@ public actor SparkSession {
3535
/// Runtime configuration interface for Spark.
3636
public let conf: RuntimeConf
3737

38+
let regexSessionID = /;session_id=([a-zA-Z0-9-]+)/
39+
3840
/// Create a session that uses the specified connection string and userID.
3941
/// - Parameters:
4042
/// - connection: a string in a patter, `sc://{host}:{port}`
4143
init(_ connection: String) {
4244
self.client = SparkConnectClient(remote: connection)
45+
// Since `Session ID` belongs to `SparkSession`, we handle this here.
46+
if connection.contains(regexSessionID) {
47+
self.sessionID = connection.firstMatch(of: regexSessionID)!.1.uppercased()
48+
} else {
49+
self.sessionID = UUID().uuidString
50+
}
4351
self.conf = RuntimeConf(self.client)
4452
}
4553

@@ -58,7 +66,7 @@ public actor SparkSession {
5866
}
5967

6068
/// A unique session ID for this session from client.
61-
nonisolated let sessionID: String = UUID().uuidString
69+
nonisolated let sessionID: String
6270

6371
/// Get the current session ID
6472
/// - Returns: the current session ID

Tests/SparkConnectTests/SparkSessionTests.swift

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,17 @@ struct SparkSessionTests {
5050
await newSpark.stop()
5151
}
5252

53+
@Test
54+
func sessionID() async throws {
55+
let spark1 = try await SparkSession.builder.getOrCreate()
56+
await spark1.stop()
57+
let remote = "sc://localhost/;session_id=\(spark1.sessionID)"
58+
let spark2 = try await SparkSession.builder.remote(remote).getOrCreate()
59+
await spark2.stop()
60+
#expect(spark1.sessionID == spark2.sessionID)
61+
#expect(spark1 == spark2)
62+
}
63+
5364
@Test func userContext() async throws {
5465
let spark = try await SparkSession.builder.getOrCreate()
5566
#if os(macOS) || os(Linux)

0 commit comments

Comments
 (0)