Scala client working

Author: Kimahriman

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala

@@ -730,6 +730,8 @@ object SparkConnectClient {
       grpcMaxMessageSize: Int = ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE,
       grpcMaxRecursionLimit: Int = ConnectCommon.CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT) {
 
+    private def isLocal = host.equals("localhost")
+
     def userContext: proto.UserContext = {
       val builder = proto.UserContext.newBuilder()
       if (userId != null) {
@@ -742,7 +744,7 @@ object SparkConnectClient {
     }
 
     def credentials: ChannelCredentials = {
-      if (isSslEnabled.contains(true)) {
+      if (isSslEnabled.contains(true) || (token.isDefined && !isLocal)) {
         token match {
           case Some(t) =>
             // With access token added in the http header.
@@ -753,22 +755,23 @@ object SparkConnectClient {
             TlsChannelCredentials.create()
         }
       } else {
-        token match {
-          case Some(t) =>
-            CompositeChannelCredentials.create(
-              InsecureChannelCredentials.create(),
-              new AccessTokenCallCredentials(t))
-          case None =>
-            InsecureChannelCredentials.create()
-        }
+        InsecureChannelCredentials.create()
       }
     }
 
     def createChannel(): ManagedChannel = {
-      val channelBuilder = Grpc.newChannelBuilderForAddress(host, port, credentials)
+      val creds = credentials
+      val channelBuilder = Grpc.newChannelBuilderForAddress(host, port, creds)
+
+      // Workaround LocalChannelCredentials are added in
+      // https://github.com/grpc/grpc-java/issues/9900
+      var metadataWithOptionalToken = metadata
+      if (!isSslEnabled.contains(true) && isLocal && token.isDefined) {
+        metadataWithOptionalToken = metadata + (("Authorization", s"Bearer ${token.get}"))
+      }
 
-      if (metadata.nonEmpty) {
-        channelBuilder.intercept(new MetadataHeaderClientInterceptor(metadata))
+      if (metadataWithOptionalToken.nonEmpty) {
+        channelBuilder.intercept(new MetadataHeaderClientInterceptor(metadataWithOptionalToken))
       }
 
       interceptors.foreach(channelBuilder.intercept(_))

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

@@ -0,0 +1,46 @@
+/*
+ * 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 io.grpc.StatusRuntimeException
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.connect.{SparkConnectServerTest, SparkSession}
+import org.apache.spark.sql.connect.service.SparkConnectService
+
+class SparkConnectAuthSuite extends SparkConnectServerTest {
+  override protected def sparkConf = {
+    super.sparkConf.set("spark.connect.authenticate.token", "deadbeef")
+  }
+
+  test("Test local authentication") {
+    val session = SparkSession
+      .builder()
+      .remote(s"sc://localhost:${SparkConnectService.localPort}/;token=deadbeef")
+      .create()
+    session.range(5).collect()
+
+    val invalidSession = SparkSession
+      .builder()
+      .remote(s"sc://localhost:${SparkConnectService.localPort}/;token=invalid")
+      .create()
+    val exception = intercept[SparkException] {
+      invalidSession.range(5).collect()
+    }
+    assert(exception.getMessage.contains("Invalid authentication token"))
+  }
+}