Merge pull request #1303 from datastax/SPARKC-635-2.5

SPARKC-635 fix per cluster settings with custom props

Author: jtgrabowski

connector/src/it/scala/com/datastax/spark/connector/sql/CassandraDataFrameSpec.scala

@@ -1,15 +1,14 @@
 package com.datastax.spark.connector.sql
 
 import java.io.IOException
-import java.util.concurrent.CompletableFuture
-
-import com.datastax.oss.driver.api.core.{CqlIdentifier, DefaultProtocolVersion}
+import com.datastax.oss.driver.api.core.{CqlIdentifier, CqlSession, DefaultProtocolVersion}
 import com.datastax.oss.driver.api.core.`type`.DataTypes
 import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder
 import com.datastax.spark.connector.cluster.DefaultCluster
 import com.datastax.spark.connector.{SparkCassandraITFlatSpecBase, _}
-import com.datastax.spark.connector.cql.{CassandraConnector, ClusteringColumn}
+import com.datastax.spark.connector.cql.{AuthConf, AuthConfFactory, CassandraConnectionFactory, CassandraConnector, CassandraConnectorConf, ClusteringColumn, DefaultConnectionFactory, NoAuthConf}
 import com.datastax.spark.connector.util.DriverUtil.toName
+import org.apache.spark.SparkConf
 import org.apache.spark.sql.SaveMode
 import org.apache.spark.sql.cassandra._
 import org.apache.spark.sql.functions._
@@ -382,4 +381,39 @@ class CassandraDataFrameSpec extends SparkCassandraITFlatSpecBase with DefaultCl
     }
   }
 
+  it should "allow to specify custom per-cluster settings" in {
+    sparkSession.conf.set("myCluster/spark.cassandra.connection.factory",
+      TestConnectionFactory.getClass.getName.split("\\$").last)
+    sparkSession.conf.set("myCluster/spark.cassandra.auth.conf.factory",
+      TestAuthFactory.getClass.getName.split("\\$").last)
+    sparkSession.conf.set("myCluster/spark.cassandra.test.custom.property", "specialValue")
+
+    sparkSession
+      .read
+      .format("org.apache.spark.sql.cassandra")
+      .options(Map("table" -> "tuple_test1", "keyspace" -> ks, "cluster" -> "myCluster"))
+      .load
+      .count()
+
+    withClue("Test auth factory was not used during the test") {
+      TestAuthFactory.used shouldBe true
+    }
+  }
+}
+
+object TestConnectionFactory extends CassandraConnectionFactory {
+  override def createSession(conf: CassandraConnectorConf): CqlSession =
+    DefaultConnectionFactory.createSession(conf)
+
+  override def properties: Set[String] = Set("spark.cassandra.test.custom.property")
 }
+
+object TestAuthFactory extends AuthConfFactory {
+  var used: Boolean = false
+
+  override def authConf(conf: SparkConf): AuthConf = {
+    used = true
+    assert(conf.get("spark.cassandra.test.custom.property") == "specialValue")
+    NoAuthConf
+  }
+}

connector/src/main/scala/com/datastax/spark/connector/cql/CassandraConnectionFactory.scala

@@ -224,11 +224,12 @@ object CassandraConnectionFactory {
     deprecatedSince = "DSE 6.0.0"
   )
 
-
   def fromSparkConf(conf: SparkConf): CassandraConnectionFactory = {
-    conf.getOption(FactoryParam.name)
-      .map(ReflectionUtil.findGlobalObject[CassandraConnectionFactory])
-      .getOrElse(FactoryParam.default)
+    fromNameOrDefault(conf.getOption(FactoryParam.name))
   }
 
+  def fromNameOrDefault(factoryName: Option[String]): CassandraConnectionFactory = {
+    factoryName.map(ReflectionUtil.findGlobalObject[CassandraConnectionFactory])
+      .getOrElse(FactoryParam.default)
+  }
 }

connector/src/main/scala/org/apache/spark/sql/cassandra/CassandraSourceRelation.scala

@@ -2,7 +2,6 @@ package org.apache.spark.sql.cassandra
 
 import java.net.InetAddress
 import java.util.{Locale, UUID}
-
 import scala.collection.mutable.ListBuffer
 import scala.util.Try
 import org.apache.hadoop.hive.conf.HiveConf
@@ -17,7 +16,7 @@ import org.apache.spark.sql.execution.datasources.LogicalRelation
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
-import com.datastax.spark.connector.cql.{CassandraConnector, CassandraConnectorConf, ColumnDef, Schema, TableDef}
+import com.datastax.spark.connector.cql.{CassandraConnectionFactory, CassandraConnector, CassandraConnectorConf, ColumnDef, Schema, TableDef}
 import com.datastax.spark.connector.rdd.partitioner.DataSizeEstimates
 import com.datastax.spark.connector.rdd.partitioner.dht.TokenFactory.forSystemLocalPartitioner
 import com.datastax.spark.connector.rdd.{CassandraJoinRDD, CassandraRDD, CassandraTableScanRDD, ReadConf}
@@ -741,19 +740,31 @@ object CassandraSourceRelation extends Logging {
     tableRef: TableRef,
     tableConf: Map[String, String]) : SparkConf = {
 
-    //Default settings
+    // Default settings
     val conf = sparkConf.clone()
     val cluster = tableRef.cluster.getOrElse(defaultClusterName)
     val ks = tableRef.keyspace
-    val AllSCCConfNames = (ConfigParameter.names ++ DeprecatedConfigParameter.names)
-    //Keyspace/Cluster level settings
-    for (prop <- AllSCCConfNames) {
-      val value = Seq(
+
+    def consolidate(prop: String): Option[String] = {
+      Seq(
         tableConf.get(prop.toLowerCase(Locale.ROOT)), //tableConf is actually a caseInsensitive map so lower case keys must be used
         sqlConf.get(s"$cluster:$ks/$prop"),
         sqlConf.get(s"$cluster/$prop"),
         sqlConf.get(s"default/$prop"),
         sqlConf.get(prop)).flatten.headOption
+    }
+
+    // Custom connection factories may have a set of custom supported properties
+    val factoryName = consolidate(CassandraConnectionFactory.FactoryParam.name)
+    val customConnectionFactoryProperties = CassandraConnectionFactory.fromNameOrDefault(factoryName).properties
+
+    val AllSCCConfNames = ConfigParameter.names ++
+      DeprecatedConfigParameter.names ++
+      customConnectionFactoryProperties
+
+    //Keyspace/Cluster level settings
+    for (prop <- AllSCCConfNames) {
+      val value = consolidate(prop)
       value.foreach(conf.set(prop, _))
     }
     //Set all user properties

driver/src/test/scala/com/datastax/spark/connector/cql/MultiplexingSchemaListenerTest.scala

@@ -7,14 +7,15 @@ package com.datastax.spark.connector.cql
 
 import com.datastax.oss.driver.api.core.`type`.UserDefinedType
 import com.datastax.oss.driver.api.core.metadata.schema.{AggregateMetadata, FunctionMetadata, KeyspaceMetadata, SchemaChangeListener, TableMetadata, ViewMetadata}
+import org.scalatest.concurrent.Eventually
 import org.scalatest.{BeforeAndAfterEach, FlatSpec, Matchers}
 import org.scalatestplus.mockito.MockitoSugar
 
 import scala.collection.mutable
 import scala.concurrent.Future
 import scala.util.Random
 
-class MultiplexingSchemaListenerTest extends FlatSpec with Matchers with MockitoSugar with BeforeAndAfterEach {
+class MultiplexingSchemaListenerTest extends FlatSpec with Matchers with MockitoSugar with BeforeAndAfterEach with Eventually {
   val r = new Random()
 
   import scala.concurrent.ExecutionContext.Implicits.global
@@ -194,7 +195,6 @@ class MultiplexingSchemaListenerTest extends FlatSpec with Matchers with Mockito
   }
 
   it should "allow listeners to be added while triggering events" in {
-    var listeners = 0
     for (it <- 1 to 200) {
       Future (listener.addListener(new IncrementingSchemaListener()))
       Future (triggerAllEvents(5))
@@ -203,7 +203,9 @@ class MultiplexingSchemaListenerTest extends FlatSpec with Matchers with Mockito
       if (it % 10 == 0)
         Future (listener.clearListeners())
     }
-    actionsDone.values.sum should be > 0
+    eventually {
+      actionsDone.values.sum should be > 0
+    }
   }
 
 }