diff --git a/clickhouse-core/src/main/scala/com/clickhouse/spark/client/ClusterClient.scala b/clickhouse-core/src/main/scala/com/clickhouse/spark/client/ClusterClient.scala index be75f8242..bf2142b1a 100644 --- a/clickhouse-core/src/main/scala/com/clickhouse/spark/client/ClusterClient.scala +++ b/clickhouse-core/src/main/scala/com/clickhouse/spark/client/ClusterClient.scala @@ -23,10 +23,12 @@ import scala.collection.JavaConverters._ import scala.util.Random._ object ClusterClient { - def apply(cluster: ClusterSpec) = new ClusterClient(cluster) + def apply(cluster: ClusterSpec, queryTimeoutMs: Long = NodeClient.DEFAULT_QUERY_TIMEOUT_MS): ClusterClient = + new ClusterClient(cluster, queryTimeoutMs) } -class ClusterClient(cluster: ClusterSpec) extends AutoCloseable with Logging { +class ClusterClient(cluster: ClusterSpec, queryTimeoutMs: Long = NodeClient.DEFAULT_QUERY_TIMEOUT_MS) + extends AutoCloseable with Logging { @transient lazy val cache = new ConcurrentHashMap[(Int, Int), NodeClient] @@ -53,7 +55,7 @@ class ClusterClient(cluster: ClusterSpec) extends AutoCloseable with Logging { val replicaSpec = shardSpec.replicas.find(_.num == r).get val nodeSpec = replicaSpec.node log.info(s"Create client to $nodeSpec, shard $s replica $r") - new NodeClient(nodeSpec) + new NodeClient(nodeSpec, queryTimeoutMs) } ) } diff --git a/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodeClient.scala b/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodeClient.scala index 0220b685f..63442fa7d 100644 --- a/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodeClient.scala +++ b/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodeClient.scala @@ -42,12 +42,14 @@ import java.util.UUID import scala.util.{Failure, Success, Try} object NodeClient { - def apply(node: NodeSpec): NodeClient = new NodeClient(node) + val DEFAULT_QUERY_TIMEOUT_MS: Long = 60000L + + def apply(node: NodeSpec, queryTimeoutMs: Long = DEFAULT_QUERY_TIMEOUT_MS): NodeClient = + new NodeClient(node, queryTimeoutMs) } -class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging { - // TODO: add configurable timeout - private val timeout: Int = 60000 +class NodeClient(val nodeSpec: NodeSpec, queryTimeoutMs: Long = NodeClient.DEFAULT_QUERY_TIMEOUT_MS) + extends AutoCloseable with Logging { private lazy val userAgent: String = { val title = getClass.getPackage.getImplementationTitle @@ -199,7 +201,7 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging { querySettings.setFormat(clickHouseFormat) querySettings.setQueryId(queryId) settings.foreach { case (k, v) => querySettings.setOption(k, v) } - Try(client.query(sql, querySettings).get(timeout, TimeUnit.MILLISECONDS)) match { + Try(client.query(sql, querySettings).get(queryTimeoutMs, TimeUnit.MILLISECONDS)) match { case Success(response: QueryResponse) => Right(deserializer(response.getInputStream)) case Failure(se: ServerException) => Left(CHServerException(se.getCode, se.getMessage, Some(nodeSpec), Some(se))) case Failure(ex: Exception) => Left(CHClientException(ex.getMessage, Some(nodeSpec), Some(ex))) @@ -233,7 +235,7 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging { querySettings.setQueryId(queryId) settings.foreach { case (k, v) => querySettings.setOption(k, v) } - Try(client.query(sql, querySettings).get(timeout, TimeUnit.MILLISECONDS)) match { + Try(client.query(sql, querySettings).get(queryTimeoutMs, TimeUnit.MILLISECONDS)) match { case Success(response: QueryResponse) => response case Failure(se: ServerException) => throw CHServerException(se.getCode, se.getMessage, Some(nodeSpec), Some(se)) case Failure(ex: Exception) => throw CHClientException(ex.getMessage, Some(nodeSpec), Some(ex)) @@ -249,6 +251,6 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging { |$sql |""".stripMargin ) - def ping(timeout: Int = timeout) = - client.ping(timeout) + def ping(timeoutMs: Long = queryTimeoutMs) = + client.ping(timeoutMs) } diff --git a/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodesClient.scala b/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodesClient.scala index d5fa1fb5e..ea00a7e65 100644 --- a/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodesClient.scala +++ b/clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodesClient.scala @@ -22,10 +22,12 @@ import scala.collection.JavaConverters._ import scala.util.Random.shuffle object NodesClient { - def apply(nodes: Nodes) = new NodesClient(nodes) + def apply(nodes: Nodes, queryTimeoutMs: Long = NodeClient.DEFAULT_QUERY_TIMEOUT_MS): NodesClient = + new NodesClient(nodes, queryTimeoutMs) } -class NodesClient(nodes: Nodes) extends AutoCloseable with Logging { +class NodesClient(nodes: Nodes, queryTimeoutMs: Long = NodeClient.DEFAULT_QUERY_TIMEOUT_MS) + extends AutoCloseable with Logging { assert(nodes.nodes.nonEmpty) @transient lazy val cache = new ConcurrentHashMap[NodeSpec, NodeClient] @@ -37,7 +39,7 @@ class NodesClient(nodes: Nodes) extends AutoCloseable with Logging { nodeSpec, { nodeSpec => log.info(s"Create client of $nodeSpec") - new NodeClient(nodeSpec) + new NodeClient(nodeSpec, queryTimeoutMs) } ) } diff --git a/docs/configurations/02_sql_configurations.md b/docs/configurations/02_sql_configurations.md index 781758373..cd6cb2b29 100644 --- a/docs/configurations/02_sql_configurations.md +++ b/docs/configurations/02_sql_configurations.md @@ -16,6 +16,7 @@ license: | |Key | Default | Description | Since |--- | ------- | ----------- | ----- +spark.clickhouse.client.queryTimeout|60s|The maximum time the ClickHouse client will wait for a single query or ping operation to complete on a NodeClient. Applied as a future-handle timeout on every client.query(...) and client.ping(...) call.|0.10.1 spark.clickhouse.ignoreUnsupportedTransform|true|ClickHouse supports using complex expressions as sharding keys or partition values, e.g. `cityHash64(col_1, col_2)`, and those can not be supported by Spark now. If `true`, ignore the unsupported expressions and log a warning, otherwise fail fast w/ an exception. Note, when `spark.clickhouse.write.distributed.convertLocal` is enabled, ignoring unsupported sharding keys may corrupt the data.|0.4.0 spark.clickhouse.read.compression.codec|lz4|The codec used to decompress data for reading. Supported codecs: none, lz4.|0.5.0 spark.clickhouse.read.distributed.convertLocal|true|When reading Distributed table, read local table instead of itself. If `true`, ignore `spark.clickhouse.read.distributed.useClusterNodes`.|0.1.0 diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala index 2d8380c11..1cdbe7ea8 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala @@ -67,7 +67,7 @@ class ClickHouseCatalog extends TableCatalog this.catalogName = name this.nodeSpec = buildNodeSpec(options) this.currentDb = nodeSpec.database - this.nodeClient = NodeClient(nodeSpec) + this.nodeClient = NodeClient(nodeSpec, clientQueryTimeoutMs) this.nodeClient.syncQueryAndCheckOutputJSONEachRow("SELECT 1") diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala index 75db6f787..98248e373 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala @@ -21,7 +21,7 @@ import com.clickhouse.spark.client.NodeClient class ClickHouseCommandRunner extends ExternalCommandRunner with ClickHouseHelper { override def executeCommand(sql: String, options: CaseInsensitiveStringMap): Array[String] = - Utils.tryWithResource(NodeClient(buildNodeSpec(options))) { nodeClient => + Utils.tryWithResource(NodeClient(buildNodeSpec(options), clientQueryTimeoutMs)) { nodeClient => nodeClient.syncQueryAndCheckOutputJSONEachRow(sql).records.map(_.toString).toArray } } diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala index 1cbda8192..ba4451f1a 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala @@ -17,7 +17,9 @@ package com.clickhouse.spark import com.clickhouse.client.ClickHouseProtocol import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.node.NullNode +import org.apache.spark.sql.catalyst.SQLConfHelper import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException} +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT import org.apache.spark.sql.clickhouse.SchemaUtils import org.apache.spark.sql.connector.catalog.Identifier import org.apache.spark.sql.types.StructType @@ -32,7 +34,7 @@ import java.time.{LocalDateTime, ZoneId} import java.util.{HashMap => JHashMap} import scala.collection.JavaConverters._ -trait ClickHouseHelper extends Logging { +trait ClickHouseHelper extends SQLConfHelper with Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_DATABASE: String => Unit = (db: String) => throw NoSuchNamespaceException(db) @@ -40,6 +42,8 @@ trait ClickHouseHelper extends Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_TABLE: (String, String) => Unit = (database, table) => throw NoSuchTableException(s"$database.$table") + def clientQueryTimeoutMs: Long = conf.getConf(CLIENT_QUERY_TIMEOUT) + def unwrap(ident: Identifier): Option[(String, String)] = ident.namespace() match { case Array(database) => Some((database, ident.name())) case _ => None diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala index 592f2a830..1a6a6e651 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala @@ -68,11 +68,12 @@ case class ClickHouseTable( lazy val (localTableSpec, localTableEngineSpec): (Option[TableSpec], Option[MergeTreeFamilyEngineSpec]) = engineSpec match { - case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) - val _localTableEngineSpec = - TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] - (Some(_localTableSpec), Some(_localTableEngineSpec)) + case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) + val _localTableEngineSpec = + TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] + (Some(_localTableSpec), Some(_localTableEngineSpec)) } case _ => (None, None) } @@ -105,8 +106,9 @@ case class ClickHouseTable( ACCEPT_ANY_SCHEMA // TODO check schema and handle extra columns before writing ).asJava - override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - queryTableSchema(database, table) + override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + queryTableSchema(database, table) } /** @@ -192,7 +194,7 @@ case class ClickHouseTable( } }.mkString("(", ",", ")") - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => dropPartition(local_db, local_table, partitionExpr, Some(cluster)) @@ -235,12 +237,13 @@ case class ClickHouseTable( val partitionSpecs: Seq[PartitionSpec] = engineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) => cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => - queryPartitionSpec(local_db, local_table) + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, clientQueryTimeoutMs)) { + implicit nodeClient: NodeClient => + queryPartitionSpec(local_db, local_table) } } case _ => - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => queryPartitionSpec(database, table) } } @@ -274,7 +277,7 @@ case class ClickHouseTable( override def deleteWhere(filters: Array[Filter]): Unit = { val deleteExpr = compileFilters(AlwaysTrue :: filters.toList) - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => delete(local_db, local_table, deleteExpr, Some(cluster)) @@ -285,7 +288,7 @@ case class ClickHouseTable( } override def truncateTable(): Boolean = - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => truncateTable(local_db, local_table, Some(cluster)) diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala index e2a8e7c5e..653339775 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala @@ -97,7 +97,8 @@ class ClickHouseScanBuilder( |$groupByClause |""".stripMargin try { - _readSchema = Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + val queryTimeoutMs = scanJob.readOptions.clientQueryTimeout + _readSchema = Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => val fields = (getQueryOutputSchema(aggQuery) zip compiledSelectItems) .map { case (structField, colExpr) => structField.copy(name = colExpr) } StructType(fields) @@ -139,10 +140,12 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch val database: String = scanJob.database val table: String = scanJob.table + private val queryTimeoutMs: Long = scanJob.readOptions.clientQueryTimeout + lazy val inputPartitions: Array[ClickHouseInputPartition] = scanJob.tableEngineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) if scanJob.readOptions.convertDistributedToLocal => scanJob.cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(local_db, local_table).map { partitionSpec => ClickHouseInputPartition( scanJob.localTableSpec.get, @@ -166,7 +169,7 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch scanJob.node )) case _: TableEngineSpec => - Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(database, table).map { partitionSpec => ClickHouseInputPartition( scanJob.tableSpec, diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala index 1e680e9c6..381d1e237 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala @@ -44,7 +44,7 @@ abstract class ClickHouseReader[Record]( // val codec: ClickHouseCompression = scanJob.readOptions.compressionCodec val readSchema: StructType = scanJob.readSchema - private lazy val nodesClient = NodesClient(part.candidateNodes) + private lazy val nodesClient = NodesClient(part.candidateNodes, scanJob.readOptions.clientQueryTimeout) def nodeClient: NodeClient = nodesClient.node diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala index 9b1809111..6bfd5c37c 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala @@ -100,7 +100,8 @@ class ClickHouseBatchWrite( log.info(s"Truncating table ${writeJob.targetDatabase(false)}.${writeJob.targetTable(false)} for overwrite mode") - Utils.tryWithResource(NodeClient(writeJob.node)) { implicit nodeClient => + val queryTimeoutMs = writeJob.writeOptions.clientQueryTimeout + Utils.tryWithResource(NodeClient(writeJob.node, queryTimeoutMs)) { implicit nodeClient => writeJob.tableEngineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => val sql = s"TRUNCATE TABLE IF EXISTS `$local_db`.`$local_table` ON CLUSTER `$cluster`" diff --git a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala index 6b715a0b9..830de16cb 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala @@ -76,6 +76,8 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) .filter(_ => writeJob.writeOptions.convertDistributedToLocal) .map(expr => SafeProjection.create(Seq(expr))) + private val queryTimeoutMs: Long = writeJob.writeOptions.clientQueryTimeout + // put the node select strategy in executor side because we need to calculate shard and don't know the records // util DataWriter#write(InternalRow) invoked. protected lazy val client: Either[ClusterClient, NodeClient] = @@ -85,11 +87,11 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) val clusterSpec = writeJob.cluster.get log.info(s"Connect to cluster ${clusterSpec.name}, which has ${clusterSpec.shards.length} shards and " + s"${clusterSpec.nodes.length} nodes.") - Left(ClusterClient(clusterSpec)) + Left(ClusterClient(clusterSpec, queryTimeoutMs)) case _ => val nodeSpec = writeJob.node log.info(s"Connect to single node: $nodeSpec") - Right(NodeClient(nodeSpec)) + Right(NodeClient(nodeSpec, queryTimeoutMs)) } def nodeClient(shardNum: Option[Int]): NodeClient = client match { diff --git a/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala b/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala index 84ec9f100..482b1bee6 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala @@ -231,4 +231,14 @@ object ClickHouseSQLConf { .transform(_.toLowerCase) .createOptional + val CLIENT_QUERY_TIMEOUT: ConfigEntry[Long] = + buildConf("spark.clickhouse.client.queryTimeout") + .doc("The maximum time the ClickHouse client will wait for a single query or ping " + + "operation to complete on a NodeClient. Applied as a future-handle timeout on every " + + "client.query(...) and client.ping(...) call.") + .version("0.10.1") + .timeConf(TimeUnit.MILLISECONDS) + .checkValue(_ > 0, "`spark.clickhouse.client.queryTimeout` should be positive.") + .createWithDefaultString("60s") + } diff --git a/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala b/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala index 0675581b1..0e193f217 100644 --- a/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala +++ b/spark-3.3/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala @@ -28,6 +28,9 @@ trait SparkOptions extends SQLConfHelper with Serializable { protected def eval[T](key: String, entry: ConfigEntry[T]): T = Option(options.get(key)).map(entry.valueConverter).getOrElse(conf.getConf(entry)) + + def clientQueryTimeout: Long = + eval(CLIENT_QUERY_TIMEOUT.key, CLIENT_QUERY_TIMEOUT) } class ReadOptions(_options: JMap[String, String]) extends SparkOptions { diff --git a/spark-3.3/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala b/spark-3.3/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala index 5cc491547..e465fa0ca 100644 --- a/spark-3.3/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala +++ b/spark-3.3/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala @@ -14,6 +14,8 @@ package org.apache.spark.sql.clickhouse +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.scalatest.funsuite.AnyFunSuite import com.clickhouse.spark.ClickHouseHelper @@ -33,4 +35,15 @@ class ClickHouseHelperSuite extends AnyFunSuite with ClickHouseHelper { assert(nodeSpec.database === "testing") assert(nodeSpec.options.get("ssl") === "true") } + + test("client query timeout uses SQLConf") { + val conf = SQLConf.get + val original = conf.getConf(CLIENT_QUERY_TIMEOUT) + assert(original === 60000L) + try { + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, "1234ms") + assert(clientQueryTimeoutMs === 1234L) + } finally + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, s"${original}ms") + } } diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala index 7a81d84e1..aeb4a4d05 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala @@ -67,7 +67,7 @@ class ClickHouseCatalog extends TableCatalog this.catalogName = name this.nodeSpec = buildNodeSpec(options) this.currentDb = nodeSpec.database - this.nodeClient = NodeClient(nodeSpec) + this.nodeClient = NodeClient(nodeSpec, clientQueryTimeoutMs) this.nodeClient.syncQueryAndCheckOutputJSONEachRow("SELECT 1") diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala index 75db6f787..98248e373 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala @@ -21,7 +21,7 @@ import com.clickhouse.spark.client.NodeClient class ClickHouseCommandRunner extends ExternalCommandRunner with ClickHouseHelper { override def executeCommand(sql: String, options: CaseInsensitiveStringMap): Array[String] = - Utils.tryWithResource(NodeClient(buildNodeSpec(options))) { nodeClient => + Utils.tryWithResource(NodeClient(buildNodeSpec(options), clientQueryTimeoutMs)) { nodeClient => nodeClient.syncQueryAndCheckOutputJSONEachRow(sql).records.map(_.toString).toArray } } diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala index 25bafc272..95a768fd0 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala @@ -17,7 +17,9 @@ package com.clickhouse.spark import com.clickhouse.client.ClickHouseProtocol import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.node.NullNode +import org.apache.spark.sql.catalyst.SQLConfHelper import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException} +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT import org.apache.spark.sql.clickhouse.SchemaUtils import org.apache.spark.sql.connector.catalog.Identifier import org.apache.spark.sql.types.StructType @@ -32,7 +34,7 @@ import java.time.{LocalDateTime, ZoneId} import java.util.{HashMap => JHashMap} import scala.collection.JavaConverters._ -trait ClickHouseHelper extends Logging { +trait ClickHouseHelper extends SQLConfHelper with Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_DATABASE: String => Unit = (db: String) => throw new NoSuchNamespaceException(db) @@ -40,6 +42,8 @@ trait ClickHouseHelper extends Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_TABLE: (String, String) => Unit = (database, table) => throw new NoSuchTableException(s"$database.$table") + def clientQueryTimeoutMs: Long = conf.getConf(CLIENT_QUERY_TIMEOUT) + def unwrap(ident: Identifier): Option[(String, String)] = ident.namespace() match { case Array(database) => Some((database, ident.name())) case _ => None diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala index 79d08bb73..09c90264d 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala @@ -69,11 +69,12 @@ case class ClickHouseTable( lazy val (localTableSpec, localTableEngineSpec): (Option[TableSpec], Option[MergeTreeFamilyEngineSpec]) = engineSpec match { - case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) - val _localTableEngineSpec = - TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] - (Some(_localTableSpec), Some(_localTableEngineSpec)) + case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) + val _localTableEngineSpec = + TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] + (Some(_localTableSpec), Some(_localTableEngineSpec)) } case _ => (None, None) } @@ -106,8 +107,9 @@ case class ClickHouseTable( ACCEPT_ANY_SCHEMA // TODO check schema and handle extra columns before writing ).asJava - override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - queryTableSchema(database, table) + override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + queryTableSchema(database, table) } /** @@ -196,7 +198,7 @@ case class ClickHouseTable( } }.mkString("(", ",", ")") - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => dropPartition(local_db, local_table, partitionExpr, Some(cluster)) @@ -239,12 +241,13 @@ case class ClickHouseTable( val partitionSpecs: Seq[PartitionSpec] = engineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) => cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => - queryPartitionSpec(local_db, local_table) + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, clientQueryTimeoutMs)) { + implicit nodeClient: NodeClient => + queryPartitionSpec(local_db, local_table) } } case _ => - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => queryPartitionSpec(database, table) } } @@ -278,7 +281,7 @@ case class ClickHouseTable( override def deleteWhere(filters: Array[Filter]): Unit = { val deleteExpr = compileFilters(AlwaysTrue :: filters.toList) - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => delete(local_db, local_table, deleteExpr, Some(cluster)) @@ -289,7 +292,7 @@ case class ClickHouseTable( } override def truncateTable(): Boolean = - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => truncateTable(local_db, local_table, Some(cluster)) diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala index e2a8e7c5e..653339775 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala @@ -97,7 +97,8 @@ class ClickHouseScanBuilder( |$groupByClause |""".stripMargin try { - _readSchema = Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + val queryTimeoutMs = scanJob.readOptions.clientQueryTimeout + _readSchema = Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => val fields = (getQueryOutputSchema(aggQuery) zip compiledSelectItems) .map { case (structField, colExpr) => structField.copy(name = colExpr) } StructType(fields) @@ -139,10 +140,12 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch val database: String = scanJob.database val table: String = scanJob.table + private val queryTimeoutMs: Long = scanJob.readOptions.clientQueryTimeout + lazy val inputPartitions: Array[ClickHouseInputPartition] = scanJob.tableEngineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) if scanJob.readOptions.convertDistributedToLocal => scanJob.cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(local_db, local_table).map { partitionSpec => ClickHouseInputPartition( scanJob.localTableSpec.get, @@ -166,7 +169,7 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch scanJob.node )) case _: TableEngineSpec => - Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(database, table).map { partitionSpec => ClickHouseInputPartition( scanJob.tableSpec, diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala index 1e680e9c6..381d1e237 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala @@ -44,7 +44,7 @@ abstract class ClickHouseReader[Record]( // val codec: ClickHouseCompression = scanJob.readOptions.compressionCodec val readSchema: StructType = scanJob.readSchema - private lazy val nodesClient = NodesClient(part.candidateNodes) + private lazy val nodesClient = NodesClient(part.candidateNodes, scanJob.readOptions.clientQueryTimeout) def nodeClient: NodeClient = nodesClient.node diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala index fcd32d66e..fbcca2981 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala @@ -102,7 +102,8 @@ class ClickHouseBatchWrite( log.info(s"Truncating table ${writeJob.targetDatabase(false)}.${writeJob.targetTable(false)} for overwrite mode") - Utils.tryWithResource(NodeClient(writeJob.node)) { implicit nodeClient => + val queryTimeoutMs = writeJob.writeOptions.clientQueryTimeout + Utils.tryWithResource(NodeClient(writeJob.node, queryTimeoutMs)) { implicit nodeClient => writeJob.tableEngineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => val sql = s"TRUNCATE TABLE IF EXISTS `$local_db`.`$local_table` ON CLUSTER `$cluster`" diff --git a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala index 9fc13f5bd..c327873af 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala @@ -96,6 +96,8 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) Some(SafeProjection.create(Seq(ExprUtils.resolveTransformCatalyst(expr, Some(writeJob.tz.getId))))) } + private val queryTimeoutMs: Long = writeJob.writeOptions.clientQueryTimeout + // put the node select strategy in executor side because we need to calculate shard and don't know the records // util DataWriter#write(InternalRow) invoked. protected lazy val client: Either[ClusterClient, NodeClient] = @@ -105,11 +107,11 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) val clusterSpec = writeJob.cluster.get log.info(s"Connect to cluster ${clusterSpec.name}, which has ${clusterSpec.shards.length} shards and " + s"${clusterSpec.nodes.length} nodes.") - Left(ClusterClient(clusterSpec)) + Left(ClusterClient(clusterSpec, queryTimeoutMs)) case _ => val nodeSpec = writeJob.node log.info(s"Connect to single node: $nodeSpec") - Right(NodeClient(nodeSpec)) + Right(NodeClient(nodeSpec, queryTimeoutMs)) } def nodeClient(shardNum: Option[Int]): NodeClient = client match { diff --git a/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala b/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala index 84ec9f100..482b1bee6 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala @@ -231,4 +231,14 @@ object ClickHouseSQLConf { .transform(_.toLowerCase) .createOptional + val CLIENT_QUERY_TIMEOUT: ConfigEntry[Long] = + buildConf("spark.clickhouse.client.queryTimeout") + .doc("The maximum time the ClickHouse client will wait for a single query or ping " + + "operation to complete on a NodeClient. Applied as a future-handle timeout on every " + + "client.query(...) and client.ping(...) call.") + .version("0.10.1") + .timeConf(TimeUnit.MILLISECONDS) + .checkValue(_ > 0, "`spark.clickhouse.client.queryTimeout` should be positive.") + .createWithDefaultString("60s") + } diff --git a/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala b/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala index 0675581b1..0e193f217 100644 --- a/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala +++ b/spark-3.4/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala @@ -28,6 +28,9 @@ trait SparkOptions extends SQLConfHelper with Serializable { protected def eval[T](key: String, entry: ConfigEntry[T]): T = Option(options.get(key)).map(entry.valueConverter).getOrElse(conf.getConf(entry)) + + def clientQueryTimeout: Long = + eval(CLIENT_QUERY_TIMEOUT.key, CLIENT_QUERY_TIMEOUT) } class ReadOptions(_options: JMap[String, String]) extends SparkOptions { diff --git a/spark-3.4/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala b/spark-3.4/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala index 5cc491547..e465fa0ca 100644 --- a/spark-3.4/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala +++ b/spark-3.4/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala @@ -14,6 +14,8 @@ package org.apache.spark.sql.clickhouse +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.scalatest.funsuite.AnyFunSuite import com.clickhouse.spark.ClickHouseHelper @@ -33,4 +35,15 @@ class ClickHouseHelperSuite extends AnyFunSuite with ClickHouseHelper { assert(nodeSpec.database === "testing") assert(nodeSpec.options.get("ssl") === "true") } + + test("client query timeout uses SQLConf") { + val conf = SQLConf.get + val original = conf.getConf(CLIENT_QUERY_TIMEOUT) + assert(original === 60000L) + try { + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, "1234ms") + assert(clientQueryTimeoutMs === 1234L) + } finally + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, s"${original}ms") + } } diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala index 48c72ffa1..a94242f12 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala @@ -78,7 +78,7 @@ class ClickHouseCatalog extends TableCatalog this.catalogName = name this.nodeSpec = buildNodeSpec(options) this.currentDb = nodeSpec.database - this.nodeClient = NodeClient(nodeSpec) + this.nodeClient = NodeClient(nodeSpec, clientQueryTimeoutMs) this.nodeClient.syncQueryAndCheckOutputJSONEachRow("SELECT 1") diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala index ff4863519..69d16dc80 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala @@ -21,7 +21,7 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap class ClickHouseCommandRunner extends ExternalCommandRunner with ClickHouseHelper { override def executeCommand(sql: String, options: CaseInsensitiveStringMap): Array[String] = - Utils.tryWithResource(client.NodeClient(buildNodeSpec(options))) { nodeClient => + Utils.tryWithResource(client.NodeClient(buildNodeSpec(options), clientQueryTimeoutMs)) { nodeClient => nodeClient.syncQueryAndCheckOutputJSONEachRow(sql).records.map(_.toString).toArray } } diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala index 97e902b2c..1eb337fa7 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala @@ -18,7 +18,9 @@ import com.clickhouse.client.ClickHouseProtocol import com.clickhouse.spark.exception.CHException import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.node.NullNode +import org.apache.spark.sql.catalyst.SQLConfHelper import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException} +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT import org.apache.spark.sql.clickhouse.SchemaUtils import org.apache.spark.sql.connector.catalog.Identifier import org.apache.spark.sql.types.StructType @@ -42,7 +44,7 @@ import java.time.{LocalDateTime, ZoneId} import java.util.{HashMap => JHashMap} import scala.collection.JavaConverters._ -trait ClickHouseHelper extends Logging { +trait ClickHouseHelper extends SQLConfHelper with Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_DATABASE: String => Unit = (db: String) => throw new NoSuchNamespaceException(db) @@ -50,6 +52,8 @@ trait ClickHouseHelper extends Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_TABLE: (String, String) => Unit = (database, table) => throw new NoSuchTableException(s"$database.$table") + def clientQueryTimeoutMs: Long = conf.getConf(CLIENT_QUERY_TIMEOUT) + def unwrap(ident: Identifier): Option[(String, String)] = ident.namespace() match { case Array(database) => Some((database, ident.name())) case _ => None diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala index 9e0ee7838..7892d1449 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala @@ -79,11 +79,12 @@ case class ClickHouseTable( lazy val (localTableSpec, localTableEngineSpec): (Option[TableSpec], Option[MergeTreeFamilyEngineSpec]) = engineSpec match { - case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) - val _localTableEngineSpec = - TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] - (Some(_localTableSpec), Some(_localTableEngineSpec)) + case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) + val _localTableEngineSpec = + TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] + (Some(_localTableSpec), Some(_localTableEngineSpec)) } case _ => (None, None) } @@ -116,8 +117,9 @@ case class ClickHouseTable( ACCEPT_ANY_SCHEMA // TODO check schema and handle extra columns before writing ).asJava - override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - queryTableSchema(database, table) + override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + queryTableSchema(database, table) } /** @@ -206,7 +208,7 @@ case class ClickHouseTable( } }.mkString("(", ",", ")") - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => dropPartition(local_db, local_table, partitionExpr, Some(cluster)) @@ -249,12 +251,13 @@ case class ClickHouseTable( val partitionSpecs: Seq[PartitionSpec] = engineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) => cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => - queryPartitionSpec(local_db, local_table) + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, clientQueryTimeoutMs)) { + implicit nodeClient: NodeClient => + queryPartitionSpec(local_db, local_table) } } case _ => - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => queryPartitionSpec(database, table) } } @@ -288,7 +291,7 @@ case class ClickHouseTable( override def deleteWhere(filters: Array[Filter]): Unit = { val deleteExpr = compileFilters(AlwaysTrue :: filters.toList) - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => delete(local_db, local_table, deleteExpr, Some(cluster)) @@ -299,7 +302,7 @@ case class ClickHouseTable( } override def truncateTable(): Boolean = - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => truncateTable(local_db, local_table, Some(cluster)) diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala index d210fb0f1..0a39f8461 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala @@ -99,7 +99,8 @@ class ClickHouseScanBuilder( |$groupByClause |""".stripMargin try { - _readSchema = Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + val queryTimeoutMs = scanJob.readOptions.clientQueryTimeout + _readSchema = Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => val fields = (getQueryOutputSchema(aggQuery) zip compiledSelectItems) .map { case (structField, colExpr) => structField.copy(name = colExpr) } StructType(fields) @@ -141,10 +142,12 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch val database: String = scanJob.database val table: String = scanJob.table + private val queryTimeoutMs: Long = scanJob.readOptions.clientQueryTimeout + lazy val inputPartitions: Array[ClickHouseInputPartition] = scanJob.tableEngineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) if scanJob.readOptions.convertDistributedToLocal => scanJob.cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(local_db, local_table).map { partitionSpec => ClickHouseInputPartition( scanJob.localTableSpec.get, @@ -168,7 +171,7 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch scanJob.node )) case _: TableEngineSpec => - Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(database, table).map { partitionSpec => ClickHouseInputPartition( scanJob.tableSpec, diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala index 53246f1bd..ede0589f7 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala @@ -43,7 +43,7 @@ abstract class ClickHouseReader[Record]( val table: String = part.table.name val readSchema: StructType = scanJob.readSchema - private lazy val nodesClient = NodesClient(part.candidateNodes) + private lazy val nodesClient = NodesClient(part.candidateNodes, scanJob.readOptions.clientQueryTimeout) def nodeClient: NodeClient = nodesClient.node diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala index fcd32d66e..fbcca2981 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala @@ -102,7 +102,8 @@ class ClickHouseBatchWrite( log.info(s"Truncating table ${writeJob.targetDatabase(false)}.${writeJob.targetTable(false)} for overwrite mode") - Utils.tryWithResource(NodeClient(writeJob.node)) { implicit nodeClient => + val queryTimeoutMs = writeJob.writeOptions.clientQueryTimeout + Utils.tryWithResource(NodeClient(writeJob.node, queryTimeoutMs)) { implicit nodeClient => writeJob.tableEngineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => val sql = s"TRUNCATE TABLE IF EXISTS `$local_db`.`$local_table` ON CLUSTER `$cluster`" diff --git a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala index 9d17feeab..916644db5 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala @@ -98,6 +98,8 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) Some(SafeProjection.create(Seq(ExprUtils.resolveTransformCatalyst(expr, Some(writeJob.tz.getId))))) } + private val queryTimeoutMs: Long = writeJob.writeOptions.clientQueryTimeout + // put the node select strategy in executor side because we need to calculate shard and don't know the records // util DataWriter#write(InternalRow) invoked. protected lazy val client: Either[ClusterClient, NodeClient] = @@ -107,11 +109,11 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) val clusterSpec = writeJob.cluster.get log.info(s"Connect to cluster ${clusterSpec.name}, which has ${clusterSpec.shards.length} shards and " + s"${clusterSpec.nodes.length} nodes.") - Left(ClusterClient(clusterSpec)) + Left(ClusterClient(clusterSpec, queryTimeoutMs)) case _ => val nodeSpec = writeJob.node log.info(s"Connect to single node: $nodeSpec") - Right(NodeClient(nodeSpec)) + Right(NodeClient(nodeSpec, queryTimeoutMs)) } def nodeClient(shardNum: Option[Int]): NodeClient = client match { diff --git a/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala b/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala index 84ec9f100..482b1bee6 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala @@ -231,4 +231,14 @@ object ClickHouseSQLConf { .transform(_.toLowerCase) .createOptional + val CLIENT_QUERY_TIMEOUT: ConfigEntry[Long] = + buildConf("spark.clickhouse.client.queryTimeout") + .doc("The maximum time the ClickHouse client will wait for a single query or ping " + + "operation to complete on a NodeClient. Applied as a future-handle timeout on every " + + "client.query(...) and client.ping(...) call.") + .version("0.10.1") + .timeConf(TimeUnit.MILLISECONDS) + .checkValue(_ > 0, "`spark.clickhouse.client.queryTimeout` should be positive.") + .createWithDefaultString("60s") + } diff --git a/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala b/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala index 0675581b1..0e193f217 100644 --- a/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala +++ b/spark-3.5/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala @@ -28,6 +28,9 @@ trait SparkOptions extends SQLConfHelper with Serializable { protected def eval[T](key: String, entry: ConfigEntry[T]): T = Option(options.get(key)).map(entry.valueConverter).getOrElse(conf.getConf(entry)) + + def clientQueryTimeout: Long = + eval(CLIENT_QUERY_TIMEOUT.key, CLIENT_QUERY_TIMEOUT) } class ReadOptions(_options: JMap[String, String]) extends SparkOptions { diff --git a/spark-3.5/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala b/spark-3.5/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala index 063b500af..e9a5add42 100644 --- a/spark-3.5/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala +++ b/spark-3.5/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala @@ -15,6 +15,8 @@ package org.apache.spark.sql.clickhouse import com.clickhouse.spark.ClickHouseHelper +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.scalatest.funsuite.AnyFunSuite @@ -33,4 +35,15 @@ class ClickHouseHelperSuite extends AnyFunSuite with ClickHouseHelper { assert(nodeSpec.database === "testing") assert(nodeSpec.options.get("ssl") === "true") } + + test("client query timeout uses SQLConf") { + val conf = SQLConf.get + val original = conf.getConf(CLIENT_QUERY_TIMEOUT) + assert(original === 60000L) + try { + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, "1234ms") + assert(clientQueryTimeoutMs === 1234L) + } finally + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, s"${original}ms") + } } diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala index aa7128eb7..f60fd7efb 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCatalog.scala @@ -78,7 +78,7 @@ class ClickHouseCatalog extends TableCatalog this.catalogName = name this.nodeSpec = buildNodeSpec(options) this.currentDb = nodeSpec.database - this.nodeClient = NodeClient(nodeSpec) + this.nodeClient = NodeClient(nodeSpec, clientQueryTimeoutMs) this.nodeClient.syncQueryAndCheckOutputJSONEachRow("SELECT 1") diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala index ff4863519..69d16dc80 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseCommandRunner.scala @@ -21,7 +21,7 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap class ClickHouseCommandRunner extends ExternalCommandRunner with ClickHouseHelper { override def executeCommand(sql: String, options: CaseInsensitiveStringMap): Array[String] = - Utils.tryWithResource(client.NodeClient(buildNodeSpec(options))) { nodeClient => + Utils.tryWithResource(client.NodeClient(buildNodeSpec(options), clientQueryTimeoutMs)) { nodeClient => nodeClient.syncQueryAndCheckOutputJSONEachRow(sql).records.map(_.toString).toArray } } diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala index a5e11d6ff..f5a7b78e4 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseHelper.scala @@ -18,7 +18,9 @@ import com.clickhouse.client.ClickHouseProtocol import com.clickhouse.spark.exception.CHException import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.node.NullNode +import org.apache.spark.sql.catalyst.SQLConfHelper import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException} +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT import org.apache.spark.sql.clickhouse.SchemaUtils import org.apache.spark.sql.connector.catalog.Identifier import org.apache.spark.sql.types.StructType @@ -42,7 +44,7 @@ import java.time.{LocalDateTime, ZoneId} import java.util.{HashMap => JHashMap} import scala.collection.JavaConverters._ -trait ClickHouseHelper extends Logging { +trait ClickHouseHelper extends SQLConfHelper with Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_DATABASE: String => Unit = (db: String) => throw new NoSuchNamespaceException(Array(db)) @@ -50,6 +52,8 @@ trait ClickHouseHelper extends Logging { @volatile lazy val DEFAULT_ACTION_IF_NO_SUCH_TABLE: (String, String) => Unit = (database, table) => throw new NoSuchTableException(database, table) + def clientQueryTimeoutMs: Long = conf.getConf(CLIENT_QUERY_TIMEOUT) + def unwrap(ident: Identifier): Option[(String, String)] = ident.namespace() match { case Array(database) => Some((database, ident.name())) case _ => None diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala index 9e0ee7838..7892d1449 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/ClickHouseTable.scala @@ -79,11 +79,12 @@ case class ClickHouseTable( lazy val (localTableSpec, localTableEngineSpec): (Option[TableSpec], Option[MergeTreeFamilyEngineSpec]) = engineSpec match { - case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) - val _localTableEngineSpec = - TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] - (Some(_localTableSpec), Some(_localTableEngineSpec)) + case distSpec: DistributedEngineSpec => Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + val _localTableSpec = queryTableSpec(distSpec.local_db, distSpec.local_table) + val _localTableEngineSpec = + TableEngineUtils.resolveTableEngine(_localTableSpec).asInstanceOf[MergeTreeFamilyEngineSpec] + (Some(_localTableSpec), Some(_localTableEngineSpec)) } case _ => (None, None) } @@ -116,8 +117,9 @@ case class ClickHouseTable( ACCEPT_ANY_SCHEMA // TODO check schema and handle extra columns before writing ).asJava - override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => - queryTableSchema(database, table) + override lazy val schema: StructType = Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { + implicit nodeClient => + queryTableSchema(database, table) } /** @@ -206,7 +208,7 @@ case class ClickHouseTable( } }.mkString("(", ",", ")") - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => dropPartition(local_db, local_table, partitionExpr, Some(cluster)) @@ -249,12 +251,13 @@ case class ClickHouseTable( val partitionSpecs: Seq[PartitionSpec] = engineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) => cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => - queryPartitionSpec(local_db, local_table) + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, clientQueryTimeoutMs)) { + implicit nodeClient: NodeClient => + queryPartitionSpec(local_db, local_table) } } case _ => - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => queryPartitionSpec(database, table) } } @@ -288,7 +291,7 @@ case class ClickHouseTable( override def deleteWhere(filters: Array[Filter]): Unit = { val deleteExpr = compileFilters(AlwaysTrue :: filters.toList) - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => delete(local_db, local_table, deleteExpr, Some(cluster)) @@ -299,7 +302,7 @@ case class ClickHouseTable( } override def truncateTable(): Boolean = - Utils.tryWithResource(NodeClient(node)) { implicit nodeClient => + Utils.tryWithResource(NodeClient(node, clientQueryTimeoutMs)) { implicit nodeClient => engineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => truncateTable(local_db, local_table, Some(cluster)) diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala index d210fb0f1..0a39f8461 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseRead.scala @@ -99,7 +99,8 @@ class ClickHouseScanBuilder( |$groupByClause |""".stripMargin try { - _readSchema = Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + val queryTimeoutMs = scanJob.readOptions.clientQueryTimeout + _readSchema = Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => val fields = (getQueryOutputSchema(aggQuery) zip compiledSelectItems) .map { case (structField, colExpr) => structField.copy(name = colExpr) } StructType(fields) @@ -141,10 +142,12 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch val database: String = scanJob.database val table: String = scanJob.table + private val queryTimeoutMs: Long = scanJob.readOptions.clientQueryTimeout + lazy val inputPartitions: Array[ClickHouseInputPartition] = scanJob.tableEngineSpec match { case DistributedEngineSpec(_, _, local_db, local_table, _, _) if scanJob.readOptions.convertDistributedToLocal => scanJob.cluster.get.shards.flatMap { shardSpec => - Utils.tryWithResource(NodeClient(shardSpec.nodes.head)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(shardSpec.nodes.head, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(local_db, local_table).map { partitionSpec => ClickHouseInputPartition( scanJob.localTableSpec.get, @@ -168,7 +171,7 @@ class ClickHouseBatchScan(scanJob: ScanJobDescription) extends Scan with Batch scanJob.node )) case _: TableEngineSpec => - Utils.tryWithResource(NodeClient(scanJob.node)) { implicit nodeClient: NodeClient => + Utils.tryWithResource(NodeClient(scanJob.node, queryTimeoutMs)) { implicit nodeClient: NodeClient => queryPartitionSpec(database, table).map { partitionSpec => ClickHouseInputPartition( scanJob.tableSpec, diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala index 53246f1bd..ede0589f7 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/read/ClickHouseReader.scala @@ -43,7 +43,7 @@ abstract class ClickHouseReader[Record]( val table: String = part.table.name val readSchema: StructType = scanJob.readSchema - private lazy val nodesClient = NodesClient(part.candidateNodes) + private lazy val nodesClient = NodesClient(part.candidateNodes, scanJob.readOptions.clientQueryTimeout) def nodeClient: NodeClient = nodesClient.node diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala index fcd32d66e..fbcca2981 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWrite.scala @@ -102,7 +102,8 @@ class ClickHouseBatchWrite( log.info(s"Truncating table ${writeJob.targetDatabase(false)}.${writeJob.targetTable(false)} for overwrite mode") - Utils.tryWithResource(NodeClient(writeJob.node)) { implicit nodeClient => + val queryTimeoutMs = writeJob.writeOptions.clientQueryTimeout + Utils.tryWithResource(NodeClient(writeJob.node, queryTimeoutMs)) { implicit nodeClient => writeJob.tableEngineSpec match { case DistributedEngineSpec(_, cluster, local_db, local_table, _, _) => val sql = s"TRUNCATE TABLE IF EXISTS `$local_db`.`$local_table` ON CLUSTER `$cluster`" diff --git a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala index 9d17feeab..916644db5 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/com/clickhouse/spark/write/ClickHouseWriter.scala @@ -98,6 +98,8 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) Some(SafeProjection.create(Seq(ExprUtils.resolveTransformCatalyst(expr, Some(writeJob.tz.getId))))) } + private val queryTimeoutMs: Long = writeJob.writeOptions.clientQueryTimeout + // put the node select strategy in executor side because we need to calculate shard and don't know the records // util DataWriter#write(InternalRow) invoked. protected lazy val client: Either[ClusterClient, NodeClient] = @@ -107,11 +109,11 @@ abstract class ClickHouseWriter(writeJob: WriteJobDescription) val clusterSpec = writeJob.cluster.get log.info(s"Connect to cluster ${clusterSpec.name}, which has ${clusterSpec.shards.length} shards and " + s"${clusterSpec.nodes.length} nodes.") - Left(ClusterClient(clusterSpec)) + Left(ClusterClient(clusterSpec, queryTimeoutMs)) case _ => val nodeSpec = writeJob.node log.info(s"Connect to single node: $nodeSpec") - Right(NodeClient(nodeSpec)) + Right(NodeClient(nodeSpec, queryTimeoutMs)) } def nodeClient(shardNum: Option[Int]): NodeClient = client match { diff --git a/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala b/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala index 4d7a2f845..eaf274b49 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/ClickHouseSQLConf.scala @@ -244,4 +244,14 @@ object ClickHouseSQLConf { .transform(_.toLowerCase) .createOptional + val CLIENT_QUERY_TIMEOUT: ConfigEntry[Long] = + buildConf("spark.clickhouse.client.queryTimeout") + .doc("The maximum time the ClickHouse client will wait for a single query or ping " + + "operation to complete on a NodeClient. Applied as a future-handle timeout on every " + + "client.query(...) and client.ping(...) call.") + .version("0.10.1") + .timeConf(TimeUnit.MILLISECONDS) + .checkValue(_ > 0, "`spark.clickhouse.client.queryTimeout` should be positive.") + .createWithDefaultString("60s") + } diff --git a/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala b/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala index 0675581b1..0e193f217 100644 --- a/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala +++ b/spark-4.0/clickhouse-spark/src/main/scala/org/apache/spark/sql/clickhouse/SparkOptions.scala @@ -28,6 +28,9 @@ trait SparkOptions extends SQLConfHelper with Serializable { protected def eval[T](key: String, entry: ConfigEntry[T]): T = Option(options.get(key)).map(entry.valueConverter).getOrElse(conf.getConf(entry)) + + def clientQueryTimeout: Long = + eval(CLIENT_QUERY_TIMEOUT.key, CLIENT_QUERY_TIMEOUT) } class ReadOptions(_options: JMap[String, String]) extends SparkOptions { diff --git a/spark-4.0/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala b/spark-4.0/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala index 063b500af..e9a5add42 100644 --- a/spark-4.0/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala +++ b/spark-4.0/clickhouse-spark/src/test/scala/org/apache/spark/sql/clickhouse/ClickHouseHelperSuite.scala @@ -15,6 +15,8 @@ package org.apache.spark.sql.clickhouse import com.clickhouse.spark.ClickHouseHelper +import org.apache.spark.sql.clickhouse.ClickHouseSQLConf.CLIENT_QUERY_TIMEOUT +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.scalatest.funsuite.AnyFunSuite @@ -33,4 +35,15 @@ class ClickHouseHelperSuite extends AnyFunSuite with ClickHouseHelper { assert(nodeSpec.database === "testing") assert(nodeSpec.options.get("ssl") === "true") } + + test("client query timeout uses SQLConf") { + val conf = SQLConf.get + val original = conf.getConf(CLIENT_QUERY_TIMEOUT) + assert(original === 60000L) + try { + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, "1234ms") + assert(clientQueryTimeoutMs === 1234L) + } finally + conf.setConfString(CLIENT_QUERY_TIMEOUT.key, s"${original}ms") + } }