Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
com.crobox.clickhouse.internal.ClickHouseExecutor.scala Maven / Gradle / Ivy
package com.crobox.clickhouse.internal
import org.apache.pekko.actor.{ActorSystem, Terminated}
import org.apache.pekko.http.scaladsl.Http
import org.apache.pekko.http.scaladsl.model._
import org.apache.pekko.http.scaladsl.settings.{ClientConnectionSettings, ConnectionPoolSettings}
import org.apache.pekko.stream._
import org.apache.pekko.stream.scaladsl.{Keep, Sink, Source, SourceQueueWithComplete}
import com.crobox.clickhouse.balancing.HostBalancer
import com.crobox.clickhouse.internal.progress.QueryProgress._
import com.crobox.clickhouse.internal.progress.{QueryProgress, StreamingProgressClickhouseTransport}
import com.crobox.clickhouse.{ClickhouseExecutionException, TooManyQueriesException}
import com.typesafe.config.Config
import com.typesafe.scalalogging.LazyLogging
import scala.concurrent.{ExecutionContext, Future, Promise}
import scala.util.{Failure, Random, Success}
private[clickhouse] trait ClickHouseExecutor extends LazyLogging {
this: ClickhouseResponseParser with ClickhouseQueryBuilder =>
protected implicit val system: ActorSystem
protected implicit val executionContext: ExecutionContext
protected val hostBalancer: HostBalancer
protected val config: Config
lazy val (progressQueue, progressSource) = {
val builtSource = QueryProgress.queryProgressStream.run()
builtSource._2.runWith(Sink.ignore) //ensure we have one sink draining the progress
builtSource
}
lazy val superPoolSettings: ConnectionPoolSettings = ConnectionPoolSettings(system)
.withConnectionSettings(
ClientConnectionSettings(system).withTransport(new StreamingProgressClickhouseTransport(progressQueue))
)
private lazy val http = Http()
private lazy val pool = http.superPool[Promise[HttpResponse]](settings = superPoolSettings)
private lazy val bufferSize: Int = config.getInt("buffer-size")
private lazy val queryRetries: Int = config.getInt("retries")
private lazy val (queue, completion) = Source
.queue[(HttpRequest, Promise[HttpResponse])](bufferSize, OverflowStrategy.backpressure)
.via(pool)
.toMat(Sink.foreach {
case (Success(resp), p) => p.success(resp)
case (Failure(e), p) => p.failure(e)
})(Keep.both)
.run()
def executeRequest(query: String,
settings: QuerySettings,
entity: Option[RequestEntity] = None,
progressQueue: Option[SourceQueueWithComplete[QueryProgress]] = None): Future[String] = {
val internalQueryIdentifier = queryIdentifier
executeWithRetries(settings.retries.getOrElse(queryRetries), progressQueue, settings) { () =>
executeRequestInternal(hostBalancer.nextHost, query, internalQueryIdentifier, settings, entity, progressQueue)
}.andThen {
case _ => progressQueue.foreach(_.complete())
}
}
protected def queryIdentifier: String =
Random.alphanumeric.take(20).mkString("")
def executeRequestWithProgress(query: String,
settings: QuerySettings,
entity: Option[RequestEntity] = None): Source[QueryProgress, Future[String]] =
Source
.queue[QueryProgress](10, OverflowStrategy.dropHead)
.mapMaterializedValue(queue => executeRequest(query, settings, entity, Some(queue)))
def shutdown(): Future[Terminated] = {
queue.complete()
queue
.watchCompletion()
.flatMap(_ => completion)
.flatMap(_ => http.shutdownAllConnectionPools())
.flatMap(_ => system.terminate())
}
protected def singleRequest(request: HttpRequest, progressEnabled: Boolean): Future[HttpResponse] = {
if(progressEnabled) {
val promise = Promise[HttpResponse]()
queue.offer(request -> promise).flatMap {
case QueueOfferResult.Enqueued => promise.future
case QueueOfferResult.Dropped => Future.failed(TooManyQueriesException())
case QueueOfferResult.QueueClosed => Future.failed(new RuntimeException(s"Queue is closed"))
case QueueOfferResult.Failure(e) => Future.failed(e)
}
} else {
http.singleRequest(request)
}
}
protected def executeRequestInternal(
host: Future[Uri],
query: String,
queryIdentifier: String,
settings: QuerySettings,
entity: Option[RequestEntity] = None,
progressQueue: Option[SourceQueueWithComplete[QueryProgress]]
): Future[String] = {
progressQueue.foreach(definedProgressQueue => {
progressSource.runForeach(
progress => {
if (progress.identifier == queryIdentifier) {
definedProgressQueue.offer(progress.progress)
}
}
)
})
host.flatMap(actualHost => {
val request = toRequest(actualHost,
query,
Some(queryIdentifier),
settings.copy(
progressHeaders = settings.progressHeaders.orElse(Some(progressQueue.isDefined))
),
entity)(config)
processClickhouseResponse(singleRequest(request, progressQueue.isDefined), query, actualHost, progressQueue)
})
}
private def executeWithRetries(retries: Int,
progressQueue: Option[SourceQueueWithComplete[QueryProgress]],
settings: QuerySettings)(
request: () => Future[String]
): Future[String] =
request().recoverWith {
case clickException: ClickhouseExecutionException if !clickException.retryable =>
// TODO use more fine grained exceptions in the client and remove the match on `Exception`
Future.failed(clickException)
case e: StreamTcpException if retries > 0 =>
progressQueue.foreach(_.offer(QueryRetry(e, (queryRetries - retries) + 1)))
logger.warn(s"Stream exception, retries left: $retries", e)
executeWithRetries(retries - 1, progressQueue, settings)(request)
case e: RuntimeException
if e.getMessage.contains("The http server closed the connection unexpectedly") && retries > 0 =>
logger.warn(s"Unexpected connection closure, retries left: $retries", e)
progressQueue.foreach(_.offer(QueryRetry(e, (queryRetries - retries) + 1)))
executeWithRetries(retries - 1, progressQueue, settings)(request)
case e: Exception if settings.idempotent.contains(true) && retries > 0 =>
logger.warn(s"Query execution exception while executing idempotent query, retries left: $retries", e)
progressQueue.foreach(_.offer(QueryRetry(e, (queryRetries - retries) + 1)))
executeWithRetries(retries - 1, progressQueue, settings)(request)
}
}
object ClickHouseExecutor {}
