fixup

Author: HyukjinKwon

python/pyspark/sql/connect/client/core.py

@@ -641,8 +641,8 @@ def __init__(
             else DefaultChannelBuilder(connection, channel_options)
         )
         self._builder.set(
-            ChannelBuilder.CONNECT_LOCAL_AUTH_TOKEN_PARAM_NAME,
-            SparkConnectClient._local_auth_token)
+            ChannelBuilder.CONNECT_LOCAL_AUTH_TOKEN_PARAM_NAME, SparkConnectClient._local_auth_token
+        )
         self._user_id = None
         self._retry_policies: List[RetryPolicy] = []
 

python/pyspark/sql/tests/connect/test_connect_session.py

@@ -313,6 +313,23 @@ def test_config(self):
         self.assertEqual(self.spark.conf.get("integer"), "1")
 
 
+@unittest.skipIf(not should_test_connect, connect_requirement_message)
+class SparkConnectLocalAuthTests(unittest.TestCase):
+    def test_auth_failure(self):
+        os.environ["SPARK_CONNECT_LOCAL_AUTH_TOKEN"] = "invalid"
+        try:
+            (
+                PySparkSession.builder.appName(self.__class__.__name__)
+                .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]"))
+                .getOrCreate()
+            )
+        except PySparkException as e:
+            assert e.getCondition() == "_LEGACY_ERROR_TEMP_3303"
+        finally:
+            del os.environ["SPARK_CONNECT_LOCAL_AUTH_TOKEN"]
+        self.fail("Exception should occur.")
+
+
 if should_test_connect:
 
     class TestError(grpc.RpcError, Exception):

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/LocalAuthInterceptor.scala

@@ -30,11 +30,10 @@ class LocalAuthInterceptor(localToken: String) extends ServerInterceptor {
       call: ServerCall[ReqT, RespT],
       headers: Metadata,
       next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
-    val t = Option(
+    val token = Option(
       headers.get(Metadata.Key
         .of(ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN_PARAM_NAME, Metadata.ASCII_STRING_MARSHALLER)))
-      .map(_.substring("Bearer ".length))
-    if (t.isEmpty || t.get != localToken) {
+    if (token.isEmpty || token.get != localToken) {
       throw new SparkSecurityException(
         errorClass = "_LEGACY_ERROR_TEMP_3303",
         messageParameters = Map.empty)

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

@@ -370,9 +370,14 @@ object SparkConnectService extends Logging {
     val sparkConnectService = new SparkConnectService(debugMode)
     val protoReflectionService =
       if (debugMode) Some(ProtoReflectionService.newInstance()) else None
+    val serverToken =
+      Option(System.getenv(ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN_ENV_NAME)).orElse {
+        if (Utils.isTesting) Some(SparkEnv.get.conf.get("spark.testing.token"))
+        else None
+      }
     val configuredInterceptors =
       SparkConnectInterceptorRegistry.createConfiguredInterceptors() ++
-        ConnectCommon.localAuthToken.map(new LocalAuthInterceptor(_))
+        serverToken.map(new LocalAuthInterceptor(_))
 
     val startServiceFn = (port: Int) => {
       val sb = bindAddress match {

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectLocalAuthE2ESuite.scala

@@ -0,0 +1,34 @@
+/*
+ * 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.spark.sql.connect.service
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.connect.SparkConnectServerTest
+
+class SparkConnectLocalAuthE2ESuite extends SparkConnectServerTest {
+  override def beforeAll(): Unit = {
+    spark.sparkContext.conf.set("spark.testing.token", "invalid")
+    super.beforeAll()
+  }
+
+  test("Test local authentication") {
+    val e = intercept[SparkException] {
+      withClient { _ => () }
+    }
+    e.getCondition == "_LEGACY_ERROR_TEMP_3303"
+  }
+}