Update java client version

.github/workflows/build-and-test.yml

@@ -13,16 +13,7 @@
 #
 
 name: "Build and Test"
-on:
-  push:
-    branches:
-      - "branch-*"
-      - "main"
-  pull_request:
-    branches:
-      - "branch-*"
-      - "main"
-  workflow_dispatch:
+on: [push]
 
 jobs:
   run-tests:

.github/workflows/cloud.yml

@@ -28,6 +28,8 @@ on:
 jobs:
   run-tests-with-clickhouse-cloud:
     runs-on: ubuntu-22.04
+    # Only run on main repository where secrets are available
+    if: github.repository == 'ClickHouse/spark-clickhouse-connector'
     strategy:
       max-parallel: 1
       fail-fast: false
@@ -44,6 +46,31 @@ jobs:
           distribution: zulu
           java-version: 8
           cache: gradle
+      - name: Wake up ClickHouse Cloud instance
+        env:
+          CLICKHOUSE_PASSWORD: ${{ secrets.INTEGRATIONS_TEAM_TESTS_CLOUD_PASSWORD_SMT }}
+        run: |
+          echo "Waking up ClickHouse Cloud instance..."
+          max_attempts=3
+          attempt=1
+
+          while [ $attempt -le $max_attempts ]; do
+            echo "Attempt $attempt of $max_attempts"
+            if curl -sS "https://${CLICKHOUSE_CLOUD_HOST}:8443/?query=SELECT+1" \
+                 --user "default:${CLICKHOUSE_PASSWORD}" \
+                 --max-time 60 > /dev/null; then
+              echo "Instance is awake!"
+              break
+            else
+              if [ $attempt -eq $max_attempts ]; then
+                echo "Failed to wake instance after $max_attempts attempts"
+                exit 1
+              fi
+              echo "Retrying in 10 seconds..."
+              sleep 10
+              ((attempt++))
+            fi
+          done
       - run: >-
           ./gradlew clean cloudTest --no-daemon --refresh-dependencies
           -Dspark_binary_version=${{ matrix.spark }}

build.gradle

@@ -90,7 +90,16 @@ allprojects {
     version = getProjectVersion()
 
     repositories {
-        maven { url = "$mavenCentralMirror" }
+        maven {
+            url = "$mavenCentralMirror"
+        }
+
+        maven {
+            url = "$mavenSnapshotsRepo"
+            mavenContent {
+                snapshotsOnly()
+            }
+        }
     }
 }
 
@@ -218,7 +227,7 @@ project(':clickhouse-core') {
         api "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:$jackson_version"
         api "com.fasterxml.jackson.module:jackson-module-scala_$scala_binary_version:$jackson_version"
 
-        api("com.clickhouse:clickhouse-jdbc:$clickhouse_jdbc_version:all") { transitive = false }
+        api("com.clickhouse:client-v2:${clickhouse_client_v2_version}:all") { transitive = false }
 
         compileOnly "jakarta.annotation:jakarta.annotation-api:$jakarta_annotation_api_version"
 
@@ -239,6 +248,7 @@ project(":clickhouse-core-it") {
         testImplementation(testFixtures(project(":clickhouse-core")))
 
         testImplementation("com.clickhouse:clickhouse-jdbc:$clickhouse_jdbc_version:all") { transitive = false }
+
         testImplementation "org.slf4j:slf4j-log4j12:$slf4j_version"
     }
 

clickhouse-core/src/main/scala/com/clickhouse/spark/client/NodeClient.scala

@@ -14,10 +14,16 @@
 
 package com.clickhouse.spark.client
 
-import com.clickhouse.spark.Logging
 import com.clickhouse.client._
-import com.clickhouse.client.config.ClickHouseClientOption
-import com.clickhouse.data.{ClickHouseCompression, ClickHouseFormat}
+import com.clickhouse.client.api.{Client, ServerException}
+import com.clickhouse.client.api.enums.Protocol
+import com.clickhouse.client.api.insert.{InsertResponse, InsertSettings}
+import com.clickhouse.client.api.query.{QueryResponse, QuerySettings}
+import com.clickhouse.data.ClickHouseFormat
+import com.clickhouse.shaded.org.apache.commons.io.IOUtils
+import com.clickhouse.spark.Logging
+
+import java.util.concurrent.TimeUnit
 import com.clickhouse.spark.exception.{CHClientException, CHException, CHServerException}
 import com.clickhouse.spark.format.{
   JSONCompactEachRowWithNamesAndTypesSimpleOutput,
@@ -30,7 +36,8 @@ import com.clickhouse.spark.spec.NodeSpec
 import com.fasterxml.jackson.databind.JsonNode
 import com.fasterxml.jackson.databind.node.ObjectNode
 
-import java.io.InputStream
+import java.io.{ByteArrayInputStream, InputStream}
+import java.time.temporal.ChronoUnit
 import java.util.UUID
 import scala.util.{Failure, Success, Try}
 
@@ -40,7 +47,7 @@ object NodeClient {
 
 class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
   // TODO: add configurable timeout
-  private val timeout: Int = 30000
+  private val timeout: Int = 60000
 
   private lazy val userAgent: String = {
     val title = getClass.getPackage.getImplementationTitle
@@ -78,24 +85,26 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
   private def shouldInferRuntime(): Boolean =
     nodeSpec.infer_runtime_env.equalsIgnoreCase("true") || nodeSpec.infer_runtime_env == "1"
 
-  private val node: ClickHouseNode = ClickHouseNode.builder()
-    .options(nodeSpec.options)
-    .host(nodeSpec.host)
-    .port(nodeSpec.protocol, nodeSpec.port)
-    .database(nodeSpec.database)
-    .credentials(ClickHouseCredentials.fromUserAndPassword(nodeSpec.username, nodeSpec.password))
-    .build()
+  private def createClickHouseURL(nodeSpec: NodeSpec): String = {
+    val ssl: Boolean = nodeSpec.options.getOrDefault("ssl", "false").toBoolean
+    if (ssl) {
+      s"https://${nodeSpec.host}:${nodeSpec.port}"
+    } else {
+      s"http://${nodeSpec.host}:${nodeSpec.port}"
+    }
+  }
 
-  private val client: ClickHouseClient = ClickHouseClient.builder()
-    .option(ClickHouseClientOption.FORMAT, ClickHouseFormat.RowBinary)
-    .option(
-      ClickHouseClientOption.PRODUCT_NAME,
-      userAgent
-    )
-    .nodeSelector(ClickHouseNodeSelector.of(node.getProtocol))
+  private val client = new Client.Builder()
+    .setUsername(nodeSpec.username)
+    .setPassword(nodeSpec.password)
+    .setDefaultDatabase(nodeSpec.database)
+    .setOptions(nodeSpec.options)
+    .setClientName(userAgent)
+    .addEndpoint(createClickHouseURL(nodeSpec))
     .build()
 
-  override def close(): Unit = client.close()
+  override def close(): Unit =
+    client.close()
 
   private def nextQueryId(): String = UUID.randomUUID.toString
 
@@ -119,15 +128,13 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
     database: String,
     table: String,
     inputFormat: String,
-    inputCompressionType: ClickHouseCompression = ClickHouseCompression.NONE,
     data: InputStream,
     settings: Map[String, String] = Map.empty
   ): Either[CHException, SimpleOutput[ObjectNode]] =
     syncInsert(
       database,
       table,
       inputFormat,
-      inputCompressionType,
       data,
       "JSONEachRow",
       JSONEachRowSimpleOutput.deserialize,
@@ -149,24 +156,32 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
     database: String,
     table: String,
     inputFormat: String,
-    inputCompressionType: ClickHouseCompression,
     data: InputStream,
     outputFormat: String,
     deserializer: InputStream => SimpleOutput[OUT],
     settings: Map[String, String]
   ): Either[CHException, SimpleOutput[OUT]] = {
+    def readAllBytes(inputStream: InputStream): Array[Byte] =
+      IOUtils.toByteArray(inputStream)
     val queryId = nextQueryId()
     val sql = s"INSERT INTO `$database`.`$table` FORMAT $inputFormat"
     onExecuteQuery(queryId, sql)
-    val req = client.write(node)
-      .query(sql, queryId)
-      .decompressClientRequest(inputCompressionType)
-      .format(ClickHouseFormat.valueOf(outputFormat))
-    settings.foreach { case (k, v) => req.set(k, v) }
-    Try(req.data(data).executeAndWait()) match {
-      case Success(resp) => Right(deserializer(resp.getInputStream))
-      case Failure(ex: ClickHouseException) =>
-        Left(CHServerException(ex.getErrorCode, ex.getMessage, Some(nodeSpec), Some(ex)))
+    val insertSettings: InsertSettings = new InsertSettings();
+    settings.foreach { case (k, v) => insertSettings.setOption(k, v) }
+    insertSettings.setDatabase(database)
+    // TODO: check what type of compression is supported by the client v2
+    insertSettings.compressClientRequest(true)
+    val payload: Array[Byte] = readAllBytes(data)
+    val is: InputStream = new ByteArrayInputStream("".getBytes())
+    Try(client.insert(
+      table,
+      new ByteArrayInputStream(payload),
+      ClickHouseFormat.valueOf(inputFormat),
+      insertSettings
+    ).get()) match {
+      case Success(resp: InsertResponse) => Right(deserializer(is))
+      case Failure(se: ServerException) =>
+        Left(CHServerException(se.getCode, se.getMessage, Some(nodeSpec), Some(se)))
       case Failure(ex) => Left(CHClientException(ex.getMessage, Some(nodeSpec), Some(ex)))
     }
   }
@@ -179,16 +194,15 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
   ): Either[CHException, SimpleOutput[OUT]] = {
     val queryId = nextQueryId()
     onExecuteQuery(queryId, sql)
-    val req = client.read(node)
-      .query(sql, queryId).asInstanceOf[ClickHouseRequest[_]]
-      .format(ClickHouseFormat.valueOf(outputFormat)).asInstanceOf[ClickHouseRequest[_]]
-      .option(ClickHouseClientOption.CONNECTION_TIMEOUT, timeout).asInstanceOf[ClickHouseRequest[_]]
-    settings.foreach { case (k, v) => req.set(k, v).asInstanceOf[ClickHouseRequest[_]] }
-    Try(req.executeAndWait()) match {
-      case Success(resp) => Right(deserializer(resp.getInputStream))
-      case Failure(ex: ClickHouseException) =>
-        Left(CHServerException(ex.getErrorCode, ex.getMessage, Some(nodeSpec), Some(ex)))
-      case Failure(ex) => Left(CHClientException(ex.getMessage, Some(nodeSpec), Some(ex)))
+    val querySettings: QuerySettings = new QuerySettings()
+    val clickHouseFormat = ClickHouseFormat.valueOf(outputFormat)
+    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 {
+      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)))
     }
   }
 
@@ -203,28 +217,26 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
   }
 
   // //////////////////////////////////////////////////////////////////////////////
-  // ///////////////////////// ret ClickHouseResponse /////////////////////////////
+  // ///////////////////////// ret QueryResponse /////////////////////////////
   // //////////////////////////////////////////////////////////////////////////////
 
   def queryAndCheck(
     sql: String,
     outputFormat: String,
-    outputCompressionType: ClickHouseCompression,
     settings: Map[String, String] = Map.empty
-  ): ClickHouseResponse = {
+  ): QueryResponse = {
     val queryId = nextQueryId()
     onExecuteQuery(queryId, sql)
-    val req = client.read(node)
-      .query(sql, queryId).asInstanceOf[ClickHouseRequest[_]]
-      .compressServerResponse(outputCompressionType).asInstanceOf[ClickHouseRequest[_]]
-      .format(ClickHouseFormat.valueOf(outputFormat)).asInstanceOf[ClickHouseRequest[_]]
-      .option(ClickHouseClientOption.CONNECTION_TIMEOUT, timeout).asInstanceOf[ClickHouseRequest[_]]
-    settings.foreach { case (k, v) => req.set(k, v).asInstanceOf[ClickHouseRequest[_]] }
-    Try(req.executeAndWait()) match {
-      case Success(resp) => resp
-      case Failure(ex: ClickHouseException) =>
-        throw CHServerException(ex.getErrorCode, ex.getMessage, Some(nodeSpec), Some(ex))
-      case Failure(ex) => throw CHClientException(ex.getMessage, Some(nodeSpec), Some(ex))
+    val querySettings: QuerySettings = new QuerySettings()
+    val clickHouseFormat = ClickHouseFormat.valueOf(outputFormat)
+    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 {
+      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))
     }
   }
 
@@ -238,5 +250,5 @@ class NodeClient(val nodeSpec: NodeSpec) extends AutoCloseable with Logging {
        |""".stripMargin
   )
   def ping(timeout: Int = timeout) =
-    client.ping(node, timeout)
+    client.ping(timeout)
 }

gradle.properties

@@ -13,7 +13,7 @@
 #
 
 mavenCentralMirror=https://repo1.maven.org/maven2/
-mavenSnapshotsRepo=https://s01.oss.sonatype.org/content/repositories/snapshots/
+mavenSnapshotsRepo=https://central.sonatype.com/repository/maven-snapshots/
 mavenReleasesRepo=https://s01.oss.sonatype.org/service/local/staging/deploy/maven2/
 
 systemProp.scala_binary_version=2.12
@@ -23,7 +23,8 @@ systemProp.known_spark_binary_versions=3.3,3.4,3.5
 
 group=com.clickhouse.spark
 
-clickhouse_jdbc_version=0.6.3
+clickhouse_jdbc_version=0.9.4
+clickhouse_client_v2_version=0.9.4
 
 spark_33_version=3.3.4
 spark_34_version=3.4.2

spark-3.3/clickhouse-spark-it/src/test/scala/org/apache/spark/sql/clickhouse/cluster/TPCDSClusterSuite.scala

@@ -38,6 +38,8 @@ class TPCDSClusterSuite extends SparkClickHouseClusterTest {
       spark.sql("CREATE DATABASE tpcds_sf1_cluster WITH DBPROPERTIES (cluster = 'single_replica')")
 
       TPCDSTestUtils.tablePrimaryKeys.foreach { case (table, primaryKeys) =>
+        println(s"before table ${table} ${primaryKeys}")
+        val start: Long = System.currentTimeMillis()
         spark.sql(
           s"""
              |CREATE TABLE tpcds_sf1_cluster.$table
@@ -51,6 +53,7 @@ class TPCDSClusterSuite extends SparkClickHouseClusterTest {
              |SELECT * FROM tpcds.sf1.$table;
              |""".stripMargin
         )
+        println(s"time took  table ${table} ${System.currentTimeMillis() - start}")
       }
 
       TPCDSTestUtils.tablePrimaryKeys.keys.foreach { table =>

spark-3.3/clickhouse-spark-it/src/test/scala/org/apache/spark/sql/clickhouse/single/ClickHouseArrowWriterSuite.scala

@@ -0,0 +1,28 @@
+/*
+ * Licensed 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
+ *
+ *     https://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.clickhouse.single
+
+import com.clickhouse.spark.base.ClickHouseSingleMixIn
+import org.apache.spark.SparkConf
+
+class ClickHouseSingleArrowWriterSuite extends ClickHouseArrowWriterSuite with ClickHouseSingleMixIn
+
+abstract class ClickHouseArrowWriterSuite extends ClickHouseWriterTestBase {
+
+  override protected def sparkConf: SparkConf = super.sparkConf
+    .set("spark.clickhouse.write.format", "arrow")
+    .set("spark.clickhouse.read.format", "json")
+
+}