This is an automated email from the ASF dual-hosted git repository. aiceflower pushed a commit to branch release-0.11.0 in repository https://gitbox.apache.org/repos/asf/linkis.git
commit 22fa84fcef3f545c0c68479c8bae124fb04fed68 Author: Wang Zhen <[email protected]> AuthorDate: Thu Feb 25 18:48:31 2021 +0800 [Linkis-565] fixbug: Presto Entrance cannot change the schema --- .../executor/PrestoEntranceEngineExecutor.scala | 53 +++++++++++++++++++++- .../PrestoEntranceEngineExecutorManager.scala | 6 +-- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutor.scala b/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutor.scala index 99a5776b9f..0ea6971275 100644 --- a/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutor.scala +++ b/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutor.scala @@ -16,9 +16,12 @@ package com.webank.wedatasphere.linkis.entrance.executor import java.sql.SQLException +import java.util import java.util.Objects +import java.util.concurrent.atomic.AtomicReference import com.facebook.presto.client.{ClientSession, QueryStatusInfo, StatementClient, StatementClientFactory} +import com.facebook.presto.spi.security.SelectedRole import com.webank.wedatasphere.linkis.common.ServiceInstance import com.webank.wedatasphere.linkis.common.io.FsPath import com.webank.wedatasphere.linkis.common.log.LogUtils @@ -46,7 +49,7 @@ import scala.collection.mutable /** * Created by yogafire on 2020/4/30 */ -class PrestoEntranceEngineExecutor(id: Long, job: EntranceJob, clientSession: ClientSession, okHttpClient: OkHttpClient, releaseResource: () => Unit) extends EntranceEngine(id = id) with SingleTaskOperateSupport with SingleTaskInfoSupport with Logging { +class PrestoEntranceEngineExecutor(id: Long, job: EntranceJob, clientSession: AtomicReference[ClientSession], okHttpClient: OkHttpClient, releaseResource: () => Unit) extends EntranceEngine(id = id) with SingleTaskOperateSupport with SingleTaskInfoSupport with Logging { private var statement: StatementClient = _ private val persistEngine = new EntranceResultSetEngine() //execute line number，as alias and progress line @@ -104,7 +107,7 @@ class PrestoEntranceEngineExecutor(id: Long, job: EntranceJob, clientSession: Cl val realCode = code.trim info(s"presto client begins to run psql code:\n $realCode") - statement = StatementClientFactory.newStatementClient(okHttpClient, clientSession, realCode) + statement = StatementClientFactory.newStatementClient(okHttpClient, clientSession.get(), realCode) initialStatusUpdates(statement) @@ -114,6 +117,9 @@ class PrestoEntranceEngineExecutor(id: Long, job: EntranceJob, clientSession: Cl } verifyServerError(statement) + + updateSession(statement) + response } @@ -242,4 +248,47 @@ class PrestoEntranceEngineExecutor(id: Long, job: EntranceJob, clientSession: Cl throw PrestoStateInvalidException("Presto status error. Statement is not finished.") } } + + private def updateSession(statement: StatementClient): Unit = { + var newSession = clientSession.get() + // update catalog and schema if present + if (statement.getSetCatalog.isPresent || statement.getSetSchema.isPresent) { + newSession = ClientSession.builder(newSession) + .withCatalog(statement.getSetCatalog.orElse(newSession.getCatalog)) + .withSchema(statement.getSetSchema.orElse(newSession.getSchema)) + .build + } + + // update transaction ID if necessary + if (statement.isClearTransactionId) newSession = ClientSession.stripTransactionId(newSession) + + var builder: ClientSession.Builder = ClientSession.builder(newSession) + + if (statement.getStartedTransactionId != null) builder = builder.withTransactionId(statement.getStartedTransactionId) + + // update session properties if present + if (!statement.getSetSessionProperties.isEmpty || !statement.getResetSessionProperties.isEmpty) { + val sessionProperties: util.Map[String, String] = new util.HashMap[String, String](newSession.getProperties) + sessionProperties.putAll(statement.getSetSessionProperties) + sessionProperties.keySet.removeAll(statement.getResetSessionProperties) + builder = builder.withProperties(sessionProperties) + } + + // update session roles + if (!statement.getSetRoles.isEmpty) { + val roles: util.Map[String, SelectedRole] = new util.HashMap[String, SelectedRole](newSession.getRoles) + roles.putAll(statement.getSetRoles) + builder = builder.withRoles(roles) + } + + // update prepared statements if present + if (!statement.getAddedPreparedStatements.isEmpty || !statement.getDeallocatedPreparedStatements.isEmpty) { + val preparedStatements: util.Map[String, String] = new util.HashMap[String, String](newSession.getPreparedStatements) + preparedStatements.putAll(statement.getAddedPreparedStatements) + preparedStatements.keySet.removeAll(statement.getDeallocatedPreparedStatements) + builder = builder.withPreparedStatements(preparedStatements) + } + + clientSession.set(newSession) + } } diff --git a/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutorManager.scala b/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutorManager.scala index 3c39363371..e17929fcc4 100644 --- a/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutorManager.scala +++ b/ujes/definedEngines/presto/entrance/src/main/scala/com/webank/wedatasphere/linkis/entrance/executor/PrestoEntranceEngineExecutorManager.scala @@ -18,7 +18,7 @@ package com.webank.wedatasphere.linkis.entrance.executor import java.net.URI import java.util import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} import java.util.{Collections, Locale, Optional, TimeZone} import com.facebook.presto.client.ClientSession @@ -92,14 +92,14 @@ class PrestoEntranceEngineExecutorManager(groupFactory: GroupFactory, val clientSession = getClientSession(job.getUser, configMap) val criteria = new SelectionCriteria(true, job.getUser, Optional.of(clientSession.getSource), Collections.emptySet(), new ResourceEstimates(Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()), Optional.empty()) - val memory: Long = new ByteType(configMap.asScala.filter(config => PRESTO_REQUEST_MEMORY.key.equals(config._1)).values.head).toLong + val memory: Long = new ByteType(PRESTO_REQUEST_MEMORY.getValue(configMap)).toLong val groupName = PrestoResourceUtils.getGroupName(criteria, PRESTO_RESOURCE_CONFIG_PATH.getValue(configMap)) val requestResource = new InstanceAndPrestoResource(new InstanceResource(1), new PrestoResource(memory, 1, groupName, PRESTO_URL.getValue)) rmClient.requestResource(job.getUser, job.getCreator, requestResource) match { case NotEnoughResource(reason) => throw new RMWarnException(40001, LogUtils.generateWarn(reason)) case AvailableResource(ticketId) => rmClient.resourceInited(UserResultResource(ticketId, job.getUser), requestResource) - new PrestoEntranceEngineExecutor(idGenerator.getAndIncrement(), job, clientSession, okHttpClient, () => rmClient.resourceReleased(UserResultResource(ticketId, job.getUser))) + new PrestoEntranceEngineExecutor(idGenerator.getAndIncrement(), job, new AtomicReference[ClientSession](clientSession), okHttpClient, () => rmClient.resourceReleased(UserResultResource(ticketId, job.getUser))) } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
