diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala index 16836b0a97d..02fc917fbe3 100644 --- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala @@ -33,7 +33,7 @@ class SetCurrentCatalog(session: Session, catalog: String) try { val session = trinoContext.clientSession.get var builder = ClientSession.builder(session) - builder = builder.withCatalog(catalog) + builder = builder.catalog(catalog) trinoContext.clientSession.set(builder.build()) setHasResultSet(false) } catch onError() diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala index aa4697f5f0e..b190251aa8b 100644 --- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala @@ -33,7 +33,7 @@ class SetCurrentDatabase(session: Session, database: String) try { val session = trinoContext.clientSession.get var builder = ClientSession.builder(session) - builder = builder.withSchema(database) + builder = builder.schema(database) trinoContext.clientSession.set(builder.build()) setHasResultSet(false) } catch onError() diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala index dd4222645c8..21aa921b9f2 100644 --- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala @@ -19,7 +19,7 @@ package org.apache.kyuubi.engine.trino.session import java.net.URI import java.time.ZoneId -import java.util.{Collections, Locale, Optional} +import java.util.{Locale, Optional} import java.util.concurrent.TimeUnit import scala.collection.JavaConverters._ @@ -93,27 +93,18 @@ class TrinoSessionImpl( val properties = getTrinoSessionConf(sessionConf).asJava - new ClientSession( - URI.create(connectionUrl), - username, - Optional.empty(), - "kyuubi", - Optional.empty(), - Collections.emptySet(), - null, - catalogName, - databaseName, - null, - ZoneId.systemDefault(), - Locale.getDefault, - Collections.emptyMap(), - properties, - Collections.emptyMap(), - Collections.emptyMap(), - Collections.emptyMap(), - null, - new Duration(clientRequestTimeout, TimeUnit.MILLISECONDS), - true) + ClientSession.builder() + .server(URI.create(connectionUrl)) + .principal(Optional.of(username)) + .source("kyuubi") + .catalog(catalogName) + .schema(databaseName) + .timeZone(ZoneId.systemDefault()) + .locale(Locale.getDefault) + .properties(properties) + .clientRequestTimeout(new Duration(clientRequestTimeout, TimeUnit.MILLISECONDS)) + .compressionDisabled(true) + .build() } private def createHttpClient(): OkHttpClient = { @@ -135,7 +126,8 @@ class TrinoSessionImpl( Optional.ofNullable(keystoreType.orNull), Optional.ofNullable(truststorePath.orNull), Optional.ofNullable(truststorePassword.orNull), - Optional.ofNullable(truststoreType.orNull), true) + Optional.ofNullable(truststoreType.orNull), + true) sessionConf.get(KyuubiConf.ENGINE_TRINO_CONNECTION_PASSWORD).foreach { password => require( diff --git a/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/WithTrinoContainerServer.scala b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/WithTrinoContainerServer.scala index dec0076ef99..f7c7b253a0f 100644 --- a/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/WithTrinoContainerServer.scala +++ b/externals/kyuubi-trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/WithTrinoContainerServer.scala @@ -23,12 +23,9 @@ import java.util.Locale import java.util.Optional import java.util.concurrent.TimeUnit -import scala.collection.JavaConverters._ - import com.dimafeng.testcontainers.TrinoContainer import com.dimafeng.testcontainers.scalatest.TestContainerForAll import io.airlift.units.Duration -import io.trino.client.ClientSelectedRole import io.trino.client.ClientSession import okhttp3.OkHttpClient import org.testcontainers.utility.DockerImageName @@ -39,7 +36,7 @@ import org.apache.kyuubi.config.KyuubiConf trait WithTrinoContainerServer extends KyuubiFunSuite with TestContainerForAll { final val IMAGE_VERSION = 411 - final val DOCKER_IMAGE_NAME = s"trinodb/trino:${IMAGE_VERSION}" + final val DOCKER_IMAGE_NAME = s"trinodb/trino:$IMAGE_VERSION" override val containerDef = TrinoContainer.Def(DockerImageName.parse(DOCKER_IMAGE_NAME)) @@ -56,27 +53,17 @@ trait WithTrinoContainerServer extends KyuubiFunSuite with TestContainerForAll { } } - protected def session(connectionUrl: String): ClientSession = new ClientSession( - URI.create(connectionUrl), - "kyuubi_test", - Optional.empty(), - "kyuubi", - Optional.empty(), - Set[String]().asJava, - null, - catalog, - schema, - null, - ZoneId.systemDefault(), - Locale.getDefault, - Map[String, String]().asJava, - Map[String, String]().asJava, - Map[String, String]().asJava, - Map[String, ClientSelectedRole]().asJava, - Map[String, String]().asJava, - null, - new Duration(2, TimeUnit.MINUTES), - true) + protected def session(connectionUrl: String): ClientSession = ClientSession.builder() + .server(URI.create(connectionUrl)) + .principal(Optional.of("kyuubi_test")) + .source("kyuubi") + .catalog(catalog) + .schema(schema) + .timeZone(ZoneId.systemDefault()) + .locale(Locale.getDefault) + .clientRequestTimeout(new Duration(2, TimeUnit.MINUTES)) + .compressionDisabled(true) + .build() lazy val httpClient = new OkHttpClient.Builder().build() } diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoClientApiSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoClientApiSuite.scala index 6ab12c05d27..8716513a81f 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoClientApiSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoClientApiSuite.scala @@ -123,32 +123,26 @@ class TrinoClientApiSuite extends KyuubiFunSuite with TrinoRestFrontendTestHelpe } private def createTestClientSession(connectUrl: URI): ClientSession = { - new ClientSession( - connectUrl, - "kyuubi_test", - Optional.of("test_user"), - "kyuubi", - Optional.of("test_token_tracing"), - Set[String]().asJava, - "test_client_info", - "test_catalog", - "test_schema", - null, - ZoneId.systemDefault(), - Locale.getDefault, - Collections.emptyMap(), - Map[String, String]( + ClientSession.builder() + .server(connectUrl) + .principal(Optional.of("kyuubi_test")) + .user(Optional.of("test_user")) + .source("kyuubi") + .traceToken(Optional.of("test_token_tracing")) + .clientInfo("test_client_info") + .catalog("test_catalog") + .schema("test_schema") + .timeZone(ZoneId.systemDefault()) + .locale(Locale.getDefault) + .properties(Map[String, String]( "test_property_key0" -> "test_property_value0", - "test_property_key1" -> "test_propert_value1").asJava, - Map[String, String]( + "test_property_key1" -> "test_propert_value1").asJava) + .preparedStatements(Map[String, String]( "test_statement_key0" -> "select 1", - "test_statement_key1" -> "select 2").asJava, - Collections.emptyMap(), - Collections.emptyMap(), - null, - new Duration(2, TimeUnit.MINUTES), - true) - + "test_statement_key1" -> "select 2").asJava) + .clientRequestTimeout(new Duration(2, TimeUnit.MINUTES)) + .compressionDisabled(true) + .build() } }