<commit_before>package services

import java.time.{Clock, Duration, LocalDateTime, OffsetDateTime, ZoneOffset}
import java.util.UUID

import akka.actor.ActorSystem
import filters.TokenAuthorizationFilter.AUTH_TOKEN_HEADER
import generated.Tables._
import generated.enums.Status
import javax.inject.{Inject, Singleton}
import models._
import org.jooq.exception.DataAccessException
import org.jooq.impl.DSL
import org.jooq.{Update => _, _}
import play.api.db.Database
import play.api.libs.ws.WSClient
import play.api.{Configuration, Logger, http}

import scala.collection.JavaConverters._
import scala.concurrent.duration._
import scala.concurrent.{Await, ExecutionContext, Future}

/**
  * A database backed queue
  *
  * @param db       database for queue state
  * @param logStore conversion log store
  */
@Singleton
class DBQueue @Inject()(db: Database,
                        logStore: LogStore,
                        ws: WSClient,
                        configuration: Configuration,
                        clock: Clock,
                        actorSystem: ActorSystem)(implicit executionContext: ExecutionContext) extends Queue with Dispatcher {

  private val logger = Logger(this.getClass)
  private val timeout = Duration.ofMillis(configuration.getMillis("queue.timeout"))

  actorSystem.scheduler.schedule(initialDelay = 10.seconds, interval = 1.minutes)(checkQueue)

  private def checkQueue(): Unit = {
    logger.debug("Check stale jobs")
    db.withConnection { connection =>
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val nowMinusTimeout = getOffsetDateTime(LocalDateTime.now(clock).minus(timeout))
      val nullTimestamp: OffsetDateTime = null
      selectJob(sql)
        .where(QUEUE.FINISHED.isNull
          .and(QUEUE.PROCESSING.isNotNull)
          .and(QUEUE.PROCESSING.lt(nowMinusTimeout))
        )
        .fetch(Mappers.JobMapper)
        .asScala
        .filter { job => !pingWorker(job) }
        .foreach { job =>
          logger.info(s"Return ${job.id} back to queue")
          try {
            sql
              .update(QUEUE)
              .set(QUEUE.STATUS, Status.queue)
              .set(QUEUE.PROCESSING, nullTimestamp)
              .where(QUEUE.UUID.eq(job.id))
              .execute()
          } catch {
            case e: DataAccessException =>
              logger.error(s"Failed to return stale job to queue: ${e.getMessage}", e)
          }
        }
    }
  }

  private def getOffsetDateTime(dateTime: LocalDateTime): OffsetDateTime = {
    val offset = clock.getZone.getRules.getOffset(dateTime)
    OffsetDateTime.of(dateTime, offset)
  }

  private def pingWorker(job: Job): Boolean = {
    // FIXME: Store worker ID to Job so we can support multiple workers
    WorkerStore.workers.values.headOption.map { worker =>
      val workerUri = worker.uri.resolve("api/v1/status")
      //          logger.debug(s"Check worker status: ${workerUri}")
      val req: Future[Boolean] = ws.url(workerUri.toString)
        .addHttpHeaders(AUTH_TOKEN_HEADER -> worker.token)
        .withRequestTimeout(10000.millis)
        .get()
        .map(_.status == http.Status.OK)
      val res: Boolean = Await.result(req, 10000.millis)
      return res
    }
    return false
  }

  private def selectJob(sql: DSLContext): SelectJoinStep[Record9[String, String, String, String, Status, Integer, OffsetDateTime, OffsetDateTime, OffsetDateTime]] = {
    sql
      .select(QUEUE.UUID, QUEUE.INPUT, QUEUE.OUTPUT, QUEUE.TRANSTYPE, QUEUE.STATUS, QUEUE.PRIORITY,
        QUEUE.CREATED, QUEUE.PROCESSING, QUEUE.FINISHED)
      .from(QUEUE)
  }

  override def contents(): Seq[Job] = {
    db.withConnection { connection =>
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val query = sql
        .select(QUEUE.UUID, QUEUE.INPUT, QUEUE.OUTPUT, QUEUE.TRANSTYPE, QUEUE.STATUS, QUEUE.PRIORITY,
          QUEUE.CREATED, QUEUE.PROCESSING, QUEUE.FINISHED)
        .from(QUEUE)
        .orderBy(QUEUE.CREATED.desc)
      val res = query
        .fetch(Mappers.JobMapper)
        .asScala
      res
    }
  }

  override def get(id: String): Option[Job] =
    db.withConnection { connection =>
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val query = sql
        .select(QUEUE.UUID, QUEUE.INPUT, QUEUE.OUTPUT, QUEUE.TRANSTYPE, QUEUE.STATUS, QUEUE.PRIORITY,
          QUEUE.CREATED, QUEUE.PROCESSING, QUEUE.FINISHED)
        .from(QUEUE)
        .where(QUEUE.UUID.eq(id))
      val res = query
        .fetchOne(Mappers.JobMapper)
      Option(res)
    }

  override def log(id: String, offset: Int): Option[Seq[String]] =
    logStore.get(id)

  override def add(job: Create): Job = {
    db.withConnection { connection =>
      val created = LocalDateTime.now(ZoneOffset.UTC)
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val query = sql
        .insertInto(QUEUE,
          QUEUE.UUID, QUEUE.CREATED, QUEUE.INPUT, QUEUE.OUTPUT, QUEUE.TRANSTYPE)
        .values(UUID.randomUUID().toString, OffsetDateTime.of(created, ZoneOffset.UTC),
          job.input, job.output, job.transtype)
      val res = query
        .returning(QUEUE.UUID, QUEUE.STATUS)
        .fetchOne()
      Job(res.getUuid, job.input, job.output, job.transtype, Map.empty, StatusString.parse(res.getStatus),
        job.priority.getOrElse(0), created, None, None)
    }
  }

  override def update(update: Update): Option[Job] = {
    db.withConnection { connection =><commit_msg>Use common query for job<commit_after>package services

import java.time.{Clock, Duration, LocalDateTime, OffsetDateTime, ZoneOffset}
import java.util.UUID

import akka.actor.ActorSystem
import filters.TokenAuthorizationFilter.AUTH_TOKEN_HEADER
import generated.Tables._
import generated.enums.Status
import javax.inject.{Inject, Singleton}
import models._
import org.jooq.exception.DataAccessException
import org.jooq.impl.DSL
import org.jooq.{Update => _, _}
import play.api.db.Database
import play.api.libs.ws.WSClient
import play.api.{Configuration, Logger, http}

import scala.collection.JavaConverters._
import scala.concurrent.duration._
import scala.concurrent.{Await, ExecutionContext, Future}

/**
  * A database backed queue
  *
  * @param db       database for queue state
  * @param logStore conversion log store
  */
@Singleton
class DBQueue @Inject()(db: Database,
                        logStore: LogStore,
                        ws: WSClient,
                        configuration: Configuration,
                        clock: Clock,
                        actorSystem: ActorSystem)(implicit executionContext: ExecutionContext) extends Queue with Dispatcher {

  private val logger = Logger(this.getClass)
  private val timeout = Duration.ofMillis(configuration.getMillis("queue.timeout"))

  actorSystem.scheduler.schedule(initialDelay = 10.seconds, interval = 1.minutes)(checkQueue)

  private def checkQueue(): Unit = {
    logger.debug("Check stale jobs")
    db.withConnection { connection =>
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val nowMinusTimeout = getOffsetDateTime(LocalDateTime.now(clock).minus(timeout))
      val nullTimestamp: OffsetDateTime = null
      selectJob(sql)
        .where(QUEUE.FINISHED.isNull
          .and(QUEUE.PROCESSING.isNotNull)
          .and(QUEUE.PROCESSING.lt(nowMinusTimeout))
        )
        .fetch(Mappers.JobMapper)
        .asScala
        .filter { job => !pingWorker(job) }
        .foreach { job =>
          logger.info(s"Return ${job.id} back to queue")
          try {
            sql
              .update(QUEUE)
              .set(QUEUE.STATUS, Status.queue)
              .set(QUEUE.PROCESSING, nullTimestamp)
              .where(QUEUE.UUID.eq(job.id))
              .execute()
          } catch {
            case e: DataAccessException =>
              logger.error(s"Failed to return stale job to queue: ${e.getMessage}", e)
          }
        }
    }
  }

  private def getOffsetDateTime(dateTime: LocalDateTime): OffsetDateTime = {
    val offset = clock.getZone.getRules.getOffset(dateTime)
    OffsetDateTime.of(dateTime, offset)
  }

  private def pingWorker(job: Job): Boolean = {
    // FIXME: Store worker ID to Job so we can support multiple workers
    WorkerStore.workers.values.headOption.map { worker =>
      val workerUri = worker.uri.resolve("api/v1/status")
      //          logger.debug(s"Check worker status: ${workerUri}")
      val req: Future[Boolean] = ws.url(workerUri.toString)
        .addHttpHeaders(AUTH_TOKEN_HEADER -> worker.token)
        .withRequestTimeout(10000.millis)
        .get()
        .map(_.status == http.Status.OK)
      val res: Boolean = Await.result(req, 10000.millis)
      return res
    }
    return false
  }

  private def selectJob(sql: DSLContext): SelectJoinStep[Record9[String, String, String, String, Status, Integer, OffsetDateTime, OffsetDateTime, OffsetDateTime]] = {
    sql
      .select(QUEUE.UUID, QUEUE.INPUT, QUEUE.OUTPUT, QUEUE.TRANSTYPE, QUEUE.STATUS, QUEUE.PRIORITY,
        QUEUE.CREATED, QUEUE.PROCESSING, QUEUE.FINISHED)
      .from(QUEUE)
  }

  override def contents(): Seq[Job] = {
    db.withConnection { connection =>
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val query = selectJob(sql)
        .orderBy(QUEUE.CREATED.desc)
      val res = query
        .fetch(Mappers.JobMapper)
        .asScala
      res
    }
  }

  override def get(id: String): Option[Job] =
    db.withConnection { connection =>
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val query = selectJob(sql)
        .where(QUEUE.UUID.eq(id))
      val res = query
        .fetchOne(Mappers.JobMapper)
      Option(res)
    }

  override def log(id: String, offset: Int): Option[Seq[String]] =
    logStore.get(id)

  override def add(job: Create): Job = {
    db.withConnection { connection =>
      val created = LocalDateTime.now(ZoneOffset.UTC)
      val sql = DSL.using(connection, SQLDialect.POSTGRES_9_4)
      val query = sql
        .insertInto(QUEUE,
          QUEUE.UUID, QUEUE.CREATED, QUEUE.INPUT, QUEUE.OUTPUT, QUEUE.TRANSTYPE)
        .values(UUID.randomUUID().toString, OffsetDateTime.of(created, ZoneOffset.UTC),
          job.input, job.output, job.transtype)
      val res = query
        .returning(QUEUE.UUID, QUEUE.STATUS)
        .fetchOne()
      Job(res.getUuid, job.input, job.output, job.transtype, Map.empty, StatusString.parse(res.getStatus),
        job.priority.getOrElse(0), created, None, None)
    }
  }

  override def update(update: Update): Option[Job] = {
    db.withConnection { connection =>
{'fraction_non_alphanumeric': 0.10971816382316901, 'fraction_numerical': 0.004524460363842021, 'mean_word_length': 2.739866055692633, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1278141', 'n_tokens_mistral': 3337, 'n_tokens_neox': 3224, 'n_words': 684}