[SPARK-44263][CONNECT] Custom Interceptors Support

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/SparkSession.scala

@@ -20,13 +20,11 @@ import java.io.Closeable
 import java.net.URI
 import java.util.concurrent.TimeUnit._
 import java.util.concurrent.atomic.AtomicLong
-
 import scala.collection.JavaConverters._
 import scala.reflect.runtime.universe.TypeTag
-
 import com.google.common.cache.{CacheBuilder, CacheLoader}
+import io.grpc.ClientInterceptor
 import org.apache.arrow.memory.RootAllocator
-
 import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.connect.proto
 import org.apache.spark.connect.proto.ExecutePlanResponse
@@ -629,7 +627,7 @@ object SparkSession extends Logging {
    * Create a new [[SparkSession]] based on the connect client [[Configuration]].
    */
   private[sql] def create(configuration: Configuration): SparkSession = {
-    new SparkSession(new SparkConnectClient(configuration), cleaner, planIdGenerator)
+    new SparkSession(configuration.toSparkConnectClient, cleaner, planIdGenerator)
   }
 
   /**
@@ -656,6 +654,16 @@ object SparkSession extends Logging {
       this
     }
 
+    /**
+     * Add an interceptor to be used during channel creation.
+     *
+     * Note that interceptors added last are executed first by gRPC.
+     */
+    def interceptor(interceptor: ClientInterceptor): Builder = {
+      builder.interceptor(interceptor)
+      this
+    }
+
     private[sql] def client(client: SparkConnectClient): Builder = {
       this.client = client
       this

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

@@ -35,9 +35,6 @@ private[sql] class SparkConnectClient(
     private[sql] val configuration: SparkConnectClient.Configuration,
     private val channel: ManagedChannel) {
 
-  def this(configuration: SparkConnectClient.Configuration) =
-    this(configuration, configuration.createChannel())
-
   private val userContext: UserContext = configuration.userContext
 
   private[this] val bstub = new CustomSparkConnectBlockingStub(channel, configuration.retryPolicy)
@@ -198,7 +195,7 @@ private[sql] class SparkConnectClient(
     bstub.interrupt(request)
   }
 
-  def copy(): SparkConnectClient = new SparkConnectClient(configuration)
+  def copy(): SparkConnectClient = configuration.toSparkConnectClient
 
   /**
    * Add a single artifact to the client session.
@@ -463,7 +460,18 @@ object SparkConnectClient {
       this
     }
 
-    def build(): SparkConnectClient = new SparkConnectClient(_configuration)
+    /**
+     * Add an interceptor to be used during channel creation.
+     *
+     * Note that interceptors added last are executed first by gRPC.
+     */
+    def interceptor(interceptor: ClientInterceptor): Builder = {
+      val interceptors = _configuration.interceptors ++ List(interceptor)
+      _configuration = _configuration.copy(interceptors = interceptors)
+      this
+    }
+
+    def build(): SparkConnectClient = _configuration.toSparkConnectClient
   }
 
   /**
@@ -478,7 +486,8 @@ object SparkConnectClient {
       isSslEnabled: Option[Boolean] = None,
       metadata: Map[String, String] = Map.empty,
       userAgent: String = DEFAULT_USER_AGENT,
-      retryPolicy: GrpcRetryHandler.RetryPolicy = GrpcRetryHandler.RetryPolicy()) {
+      retryPolicy: GrpcRetryHandler.RetryPolicy = GrpcRetryHandler.RetryPolicy(),
+      interceptors: List[ClientInterceptor] = List.empty) {
 
     def userContext: proto.UserContext = {
       val builder = proto.UserContext.newBuilder()
@@ -509,12 +518,18 @@ object SparkConnectClient {
 
     def createChannel(): ManagedChannel = {
       val channelBuilder = Grpc.newChannelBuilderForAddress(host, port, credentials)
+
       if (metadata.nonEmpty) {
         channelBuilder.intercept(new MetadataHeaderClientInterceptor(metadata))
       }
+
+      interceptors.forEach(channelBuilder.intercept(_))
+
       channelBuilder.maxInboundMessageSize(ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE)
       channelBuilder.build()
     }
+
+    def toSparkConnectClient: SparkConnectClient = new SparkConnectClient(this, createChannel())
   }
 
   /**
@@ -540,10 +555,6 @@ object SparkConnectClient {
         }
       })
     }
-
-    override def thisUsesUnstableApi(): Unit = {
-      // Marks this API is not stable. Left empty on purpose.
-    }
   }
 
   /**

connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/SparkSessionSuite.scala

@@ -16,12 +16,15 @@
  */
 package org.apache.spark.sql
 
+
+import io.grpc.{CallOptions, Channel, ClientCall, ClientInterceptor, MethodDescriptor, Server}
+
 import org.apache.spark.sql.connect.client.util.ConnectFunSuite
 
 /**
  * Tests for non-dataframe related SparkSession operations.
  */
-class SparkSessionSuite extends ConnectFunSuite {
+class SparkSessionSuite extends ConnectFunSuite with BeforeAndAfterEach {
   test("default") {
     val session = SparkSession.builder().getOrCreate()
     assert(session.client.configuration.host == "localhost")
@@ -73,4 +76,20 @@ class SparkSessionSuite extends ConnectFunSuite {
       session2.close()
     }
   }
+
+  test("Custom Interceptor") {
+    val session = SparkSession
+      .builder()
+      .interceptor(new ClientInterceptor {
+        override def interceptCall[ReqT, RespT](methodDescriptor: MethodDescriptor[ReqT, RespT],
+                                                callOptions: CallOptions,
+                                                channel: Channel): ClientCall[ReqT, RespT] = {
+          throw new RuntimeException("Blocked")
+        }
+      }).create()
+
+    assertThrows[RuntimeException] {
+      session.range(10).count()
+    }
+  }
 }

connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala

@@ -21,7 +21,7 @@ import java.util.concurrent.TimeUnit
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 
-import io.grpc.{Server, Status, StatusRuntimeException}
+import io.grpc.{CallOptions, Channel, ClientCall, ClientInterceptor, MethodDescriptor, Server}
 import io.grpc.netty.NettyServerBuilder
 import io.grpc.stub.StreamObserver
 import org.scalatest.BeforeAndAfterEach
@@ -123,6 +123,26 @@ class SparkConnectClientSuite extends ConnectFunSuite with BeforeAndAfterEach {
     assert(df.plan === service.getAndClearLatestInputPlan())
   }
 
+  test("Custom Interceptor") {
+    startDummyServer(0)
+    client = SparkConnectClient
+      .builder()
+      .connectionString(s"sc://localhost:${server.getPort}")
+      .interceptor(new ClientInterceptor {
+        override def interceptCall[ReqT, RespT](methodDescriptor: MethodDescriptor[ReqT, RespT],
+                                                callOptions: CallOptions,
+                                                channel: Channel): ClientCall[ReqT, RespT] = {
+          throw new RuntimeException("Blocked")
+        }
+      }).build()
+
+    val session = SparkSession.builder().client(client).create()
+
+    assertThrows[RuntimeException] {
+      session.range(10).count()
+    }
+  }
+
   private case class TestPackURI(
       connectionString: String,
       isCorrect: Boolean,