[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,11 @@ object SparkSession extends Logging {
       this
     }
 
+    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 stub = new CustomSparkConnectBlockingStub(channel)
@@ -197,7 +194,7 @@ private[sql] class SparkConnectClient(
     stub.interrupt(request)
   }
 
-  def copy(): SparkConnectClient = new SparkConnectClient(configuration)
+  def copy(): SparkConnectClient = configuration.toSparkConnectClient
 
   /**
    * Add a single artifact to the client session.
@@ -457,7 +454,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
   }
 
   /**
@@ -471,7 +479,8 @@ object SparkConnectClient {
       token: Option[String] = None,
       isSslEnabled: Option[Boolean] = None,
       metadata: Map[String, String] = Map.empty,
-      userAgent: String = DEFAULT_USER_AGENT) {
+      userAgent: String = DEFAULT_USER_AGENT,
+      interceptors: List[ClientInterceptor] = List.empty) {
 
     def userContext: proto.UserContext = {
       val builder = proto.UserContext.newBuilder()
@@ -502,12 +511,20 @@ object SparkConnectClient {
 
     def createChannel(): ManagedChannel = {
       val channelBuilder = Grpc.newChannelBuilderForAddress(host, port, credentials)
+
       if (metadata.nonEmpty) {
         channelBuilder.intercept(new MetadataHeaderClientInterceptor(metadata))
       }
+
+      for (interceptor <- interceptors) {
+        channelBuilder.intercept(interceptor)
+      }
+
       channelBuilder.maxInboundMessageSize(ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE)
       channelBuilder.build()
     }
+
+    def toSparkConnectClient: SparkConnectClient = new SparkConnectClient(this, createChannel())
   }
 
   /**
@@ -533,10 +550,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/connect/client/SparkConnectClientSuite.scala

@@ -17,7 +17,6 @@
 package org.apache.spark.sql.connect.client
 
 import java.util.concurrent.TimeUnit
-
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 
@@ -123,6 +122,28 @@ class SparkConnectClientSuite extends ConnectFunSuite with BeforeAndAfterEach {
     assert(df.plan === service.getAndClearLatestInputPlan())
   }
 
+  test("CustomInterceptor") {
+    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,