fix

apache · Jan 15, 2024 · c498a5b · c498a5b
1 parent 21948ca
commit c498a5b
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 74 deletions.
diff --git a/...no-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentCatalog.scala b/...no-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/...o-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/SetCurrentDatabase.scala b/...o-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/...trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala b/...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/...trino-engine/src/test/scala/org/apache/kyuubi/engine/trino/WithTrinoContainerServer.scala b/...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
@@ -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()
   }
 
 }