<commit_before>package controllers.graphql

import akka.http.scaladsl.model.HttpResponse
import akka.http.scaladsl.model.ws.{Message, TextMessage, UpgradeToWebSocket}
import akka.stream.scaladsl.{Flow, Keep, Sink, Source, SourceQueueWithComplete}
import akka.stream.{ActorMaterializer, KillSwitches, OverflowStrategy, SharedKillSwitch}
import controllers.graphql.jsonProtocols.OperationMessage
import controllers.graphql.jsonProtocols.OperationMessageJsonProtocol._
import controllers.graphql.jsonProtocols.OperationMessageType._
import graphql.{GraphQLContext, GraphQLContextFactory}
import javax.inject.{Inject, Singleton}
import monix.execution.Scheduler
import sangria.ast.OperationType.Subscription
import sangria.execution.ExecutionScheme.Stream
import sangria.execution.Executor
import sangria.marshalling.sprayJson._
import sangria.parser.{QueryParser, SyntaxError}
import spray.json._

import scala.util.{Failure, Success}

@Singleton
class WebSocketHandler @Inject()(graphQlContextFactory: GraphQLContextFactory,
                                 graphQlExecutor: Executor[GraphQLContext, Unit])
                                (implicit val actorMaterializer: ActorMaterializer,
                                 implicit val scheduler: Scheduler) {

  private val graphqlWebsocketProtocol = Some("graphql-ws")

  def handleMessages(upgradeToWebSocket: UpgradeToWebSocket): HttpResponse = {
    implicit val (queue, publisher) = Source.queue[Message](0, OverflowStrategy.fail)
      .toMat(Sink.asPublisher(false))(Keep.both)
      .run()
    val killSwitches = KillSwitches.shared(this.getClass.getSimpleName)
    val incoming = Flow[Message]
      .collect {
        case TextMessage.Strict(query) =>
          val operation = query.parseJson.convertTo[OperationMessage]
          operation.operationType match {
            case GQL_CONNECTION_INIT =>
              reply(OperationMessage(GQL_CONNECTION_ACK, None, None))
            case GQL_START =>
              handleGraphQlQuery(operation, killSwitches)
          }
      }
      .to {
        Sink.onComplete {
          _ =>
            killSwitches.shutdown
            queue.complete
        }
      }
    upgradeToWebSocket.handleMessagesWithSinkSource(incoming, Source.fromPublisher(publisher), graphqlWebsocketProtocol)
  }

  private def handleGraphQlQuery(operationMessage: OperationMessage, killSwitches: SharedKillSwitch)
                                (implicit queue: SourceQueueWithComplete[Message]): Unit = {

    import sangria.streaming.akkaStreams._

    operationMessage.payload.foreach {
      payload =>
        val JsObject(fields) = payload
        val JsString(query) = fields("query")
        val operation = fields.get("operationName") collect {
          case JsString(op) => op
        }
        val vars = fields.get("variables") match {
          case Some(obj: JsObject) => obj
          case _ => JsObject.empty
        }
        QueryParser.parse(query) match {
          case Success(queryAst) =>
            queryAst.operationType(operation) match {
              case Some(Subscription) =>
                val ctx = graphQlContextFactory.createContextForRequest

                graphQlExecutor.execute(
                  queryAst = queryAst,
                  userContext = ctx,
                  root = (),
                  operationName = operation,
                  variables = vars
                ).viaMat(killSwitches.flow)(Keep.none)
                  .runForeach {
                    result =>
                      reply(OperationMessage(GQL_DATA, operationMessage.id, Some(result)))
                  }
              case _ =>
                reply(OperationMessage(<commit_msg>Add using GraphQLMessage for websocket handler<commit_after>package controllers.graphql

import akka.http.scaladsl.model.HttpResponse
import akka.http.scaladsl.model.ws.{Message, TextMessage, UpgradeToWebSocket}
import akka.stream.scaladsl.{Flow, Keep, Sink, Source, SourceQueueWithComplete}
import akka.stream.{ActorMaterializer, KillSwitches, OverflowStrategy, SharedKillSwitch}
import controllers.graphql.jsonProtocols.GraphQLMessageProtocol._
import controllers.graphql.jsonProtocols.OperationMessageJsonProtocol._
import controllers.graphql.jsonProtocols.OperationMessageType._
import controllers.graphql.jsonProtocols.{GraphQLMessage, OperationMessage}
import graphql.{GraphQLContext, GraphQLContextFactory}
import javax.inject.{Inject, Singleton}
import monix.execution.Scheduler
import sangria.ast.OperationType.Subscription
import sangria.execution.ExecutionScheme.Stream
import sangria.execution.Executor
import sangria.marshalling.sprayJson._
import sangria.parser.{QueryParser, SyntaxError}
import spray.json._

import scala.util.{Failure, Success}

@Singleton
class WebSocketHandler @Inject()(graphQlContextFactory: GraphQLContextFactory,
                                 graphQlExecutor: Executor[GraphQLContext, Unit])
                                (implicit val actorMaterializer: ActorMaterializer,
                                 implicit val scheduler: Scheduler) {

  private val graphqlWebsocketProtocol = Some("graphql-ws")

  def handleMessages(upgradeToWebSocket: UpgradeToWebSocket): HttpResponse = {
    implicit val (queue, publisher) = Source.queue[Message](0, OverflowStrategy.fail)
      .toMat(Sink.asPublisher(false))(Keep.both)
      .run()
    val killSwitches = KillSwitches.shared(this.getClass.getSimpleName)
    val incoming = Flow[Message]
      .collect {
        case TextMessage.Strict(query) =>
          val operation = query.parseJson.convertTo[OperationMessage]
          operation.operationType match {
            case GQL_CONNECTION_INIT =>
              reply(OperationMessage(GQL_CONNECTION_ACK, None, None))
            case GQL_START =>
              handleGraphQlQuery(operation, killSwitches)
          }
      }
      .to {
        Sink.onComplete {
          _ =>
            killSwitches.shutdown
            queue.complete
        }
      }
    upgradeToWebSocket.handleMessagesWithSinkSource(incoming, Source.fromPublisher(publisher), graphqlWebsocketProtocol)
  }

  private def handleGraphQlQuery(operationMessage: OperationMessage, killSwitches: SharedKillSwitch)
                                (implicit queue: SourceQueueWithComplete[Message]): Unit = {

    import sangria.streaming.akkaStreams._

    operationMessage.payload.foreach {
      payload =>
        val graphQlMessage = payload.convertTo[GraphQLMessage]
        QueryParser.parse(graphQlMessage.query) match {
          case Success(queryAst) =>
            queryAst.operationType(operation) match {
              case Some(Subscription) =>
                val ctx = graphQlContextFactory.createContextForRequest
                graphQlExecutor.execute(
                  queryAst = queryAst,
                  userContext = ctx,
                  root = (),
                  operationName = graphQlMessage.operationName,
                  variables = graphQlMessage.variables.getOrElse(JsObject.empty)
                ).viaMat(killSwitches.flow)(Keep.none)
                  .runForeach {
                    result =>
                      reply(OperationMessage(GQL_DATA, operationMessage.id, Some(result)))
                  }
              case _ =>
                reply(OperationMessage(
{'fraction_non_alphanumeric': 0.08026208026208026, 'fraction_numerical': 0.000273000273000273, 'mean_word_length': 2.6856136820925554, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5104677', 'n_tokens_mistral': 1941, 'n_tokens_neox': 1870, 'n_words': 341}