<commit_before>package io.waylay.kairosdb.driver

import java.io.ByteArrayOutputStream
import java.util.zip.GZIPOutputStream

import akka.util.ByteString
import io.waylay.kairosdb.driver.models._
import io.waylay.kairosdb.driver.models.HealthCheckResult._
import com.netaporter.uri.dsl._
import com.typesafe.scalalogging.StrictLogging
import io.waylay.kairosdb.driver.KairosDB._
import io.waylay.kairosdb.driver.models.QueryMetricTagsResponse.TagsResponse
import io.waylay.kairosdb.driver.models.json.Formats._
import io.waylay.kairosdb.driver.models.QueryResponse.Response
import play.api.libs.json._
import play.api.libs.ws.{StandaloneWSClient, StandaloneWSRequest, StandaloneWSResponse, WSAuthScheme}
import play.api.libs.ws.DefaultBodyWritables._
import play.api.libs.ws.JsonBodyWritables._
import play.api.libs.ws.JsonBodyReadables._

import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success, Try}
import scala.concurrent.duration._

object KairosDB {

  /** Parent exception for all exceptions this driver may throw */
  class KairosDBException(msg: String) extends Exception(msg)

  /** Parent exception for all exception this driver may throw when querying data */
  abstract class KairosDBResponseException(msg: String) extends KairosDBException(msg)

  /** Thrown when the response by KairosDB can't be parsed */
  case class KairosDBResponseParseException(msg: String = "Could not parse KairosDB response", errors: Seq[(JsPath, Seq[JsonValidationError])])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown when the request is invalid */
  case class KairosDBResponseBadRequestException(msg: String = "KairosDB returned HTTP 400 (Bad Request)", errors: Seq[String])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown when the request is unauthorized */
  case class KairosDBResponseUnauthorizedException(msg: String = "KairosDB returned HTTP 401 (Unauthorized)", errors: Seq[String])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown if an error occurs retrieving data */
  case class KairosDBResponseInternalServerErrorException(msg: String = "KairosDB returned HTTP 500 (Internal Server Error)", errors: Seq[String])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown if an unexpected error occurs retrieving data */
  case class KairosDBResponseUnhandledException(statusCode: Int, errors: Seq[String])
    extends KairosDBResponseException(s"""KairosDB returned unhandled HTTP $statusCode (???) (errors: ${errors.mkString(", ")})""")
}

class KairosDB(wsClient: StandaloneWSClient, config: KairosDBConfig, executionContext: ExecutionContext) extends StrictLogging {
  implicit val ec = executionContext
  val uri = config.uri

  def listMetricNames: Future[Seq[MetricName]] = {
    wsClient
      .url(uri / "api" / "v1" / "metricnames")
      .applyKairosDBAuth
      .get
      .map {
        wsRepsonseToResult(json => (json \ "results").validate[Seq[String]])
      }
      .map(_.map(MetricName))
  }

  def listTagNames: Future[Seq[String]] = {
    wsClient
      .url(uri / "api" / "v1" / "tagnames")
      .applyKairosDBAuth
      .get
      .map {
        wsRepsonseToResult(json => (json \ "results").validate[Seq[String]])
      }
  }

  def listTagValues: Future[Seq[String]] = {
    wsClient
      .url(uri / "api" / "v1" / "tagvalues")
      .applyKairosDBAuth
      .get
      .map(
        wsRepsonseToResult(json => (json \ "results").validate[Seq[String]])
      )
  }

  def healthStatus: Future[HealthStatusResults] = {
    wsClient
      .url(uri / "api" / "v1" / "health" / "status")
      .applyKairosDBAuth
      .get
      .map {
        wsRepsonseToResult(_.validate[Seq[String]])
      }
      .map(HealthStatusResults)
  }

  def healthCheck: Future[HealthCheckResult] = {
    wsClient
      .url(uri / "api" / "v1" / "health" / "check")
      .withRequestTimeout(10.seconds)
      .applyKairosDBAuth
      .get
      .map(_.status)
      .map {
        case 204 => AllHealthy
        case _   => Unhealthy
      }
      .recover {
        case _ => Unhealthy
      }
  }

  def version: Future[String] = {
    wsClient
      .url(uri / "api" / "v1" / "version")
      .applyKairosDBAuth
      .get
      .map (
        wsRepsonseToResult(json => (json \ "version").validate[String])
      )
  }

  def deleteMetric(metricName: MetricName): Future[Unit] = {
    wsClient
      .url(uri / "api" / "v1" / "metric" / metricName.name)
      .applyKairosDBAuth
      .delete
      .map(emptyWsRepsonseToResult)
  }

  def addDataPoints(dataPoints: Seq[DataPoint], gzip: Boolean = false): Future[Unit] = {
    val body = Json.toJson(dataPoints)
    logger.debug(Json.prettyPrint(body))

    val (actualBody, contentType) = if(gzip){
      val buff = new ByteArrayOutputStream()
      val gzos = new GZIPOutputStream(buff)
      gzos.write(Json.toBytes(body))
      gzos.finish()
      val gzipBody = buff.toByteArray
      (gzipBody, "application/gzip")
    }else{
      (Json.toBytes(body), "application/json")
    }

    wsClient
      .url(uri / "api" / "v1" / "datapoints")
      .applyKairosDBAuth
      .addHttpHeaders("Content-Type" -> contentType)
      .post(ByteString(actualBody))
      .map(emptyWsRepsonseToResult)
  }

  def addDataPoint(dataPoint: DataPoint): Future[Unit] = addDataPoints(Seq(dataPoint))

  def queryMetrics(queryMetrics: QueryMetrics): Future[Response] = {
    val query = Json.toJson(queryMetrics)
    logger.debug(Json.prettyPrint(query))
    wsClient
      .url(uri / "api" / "v1" / "datapoints" / "query")
      .applyKairosDBAuth
      .post(query)
      .map {
        wsRepsonseToResult(_.validate[Response])
      }
  }

  /** You can think of this as the exact same as the query but it leaves off the data and just returns the tag information.
    *
    * Note: Currently this is not implemented in the HBase datastore.
    */
  def queryMetricTags(queryMetrics: QueryMetrics): Future[TagsResponse] = {
    val query = Json.toJson(queryMetrics)
    logger.debug(Json.prettyPrint(query))
    wsClient
      .url(uri / "api" / "v1" / "datapoints" / "query" / "tags")<commit_msg>removes unneeded byte buffer wrap<commit_after>package io.waylay.kairosdb.driver

import java.io.ByteArrayOutputStream
import java.util.zip.GZIPOutputStream

import io.waylay.kairosdb.driver.models._
import io.waylay.kairosdb.driver.models.HealthCheckResult._
import com.netaporter.uri.dsl._
import com.typesafe.scalalogging.StrictLogging
import io.waylay.kairosdb.driver.KairosDB._
import io.waylay.kairosdb.driver.models.QueryMetricTagsResponse.TagsResponse
import io.waylay.kairosdb.driver.models.json.Formats._
import io.waylay.kairosdb.driver.models.QueryResponse.Response
import play.api.libs.json._
import play.api.libs.ws.{StandaloneWSClient, StandaloneWSRequest, StandaloneWSResponse, WSAuthScheme}
import play.api.libs.ws.DefaultBodyWritables._
import play.api.libs.ws.JsonBodyWritables._
import play.api.libs.ws.JsonBodyReadables._

import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success, Try}
import scala.concurrent.duration._

object KairosDB {

  /** Parent exception for all exceptions this driver may throw */
  class KairosDBException(msg: String) extends Exception(msg)

  /** Parent exception for all exception this driver may throw when querying data */
  abstract class KairosDBResponseException(msg: String) extends KairosDBException(msg)

  /** Thrown when the response by KairosDB can't be parsed */
  case class KairosDBResponseParseException(msg: String = "Could not parse KairosDB response", errors: Seq[(JsPath, Seq[JsonValidationError])])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown when the request is invalid */
  case class KairosDBResponseBadRequestException(msg: String = "KairosDB returned HTTP 400 (Bad Request)", errors: Seq[String])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown when the request is unauthorized */
  case class KairosDBResponseUnauthorizedException(msg: String = "KairosDB returned HTTP 401 (Unauthorized)", errors: Seq[String])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown if an error occurs retrieving data */
  case class KairosDBResponseInternalServerErrorException(msg: String = "KairosDB returned HTTP 500 (Internal Server Error)", errors: Seq[String])
    extends KairosDBResponseException(s"""$msg (errors: ${errors.mkString(", ")})""")

  /** Thrown if an unexpected error occurs retrieving data */
  case class KairosDBResponseUnhandledException(statusCode: Int, errors: Seq[String])
    extends KairosDBResponseException(s"""KairosDB returned unhandled HTTP $statusCode (???) (errors: ${errors.mkString(", ")})""")
}

class KairosDB(wsClient: StandaloneWSClient, config: KairosDBConfig, executionContext: ExecutionContext) extends StrictLogging {
  implicit val ec = executionContext
  val uri = config.uri

  def listMetricNames: Future[Seq[MetricName]] = {
    wsClient
      .url(uri / "api" / "v1" / "metricnames")
      .applyKairosDBAuth
      .get
      .map {
        wsRepsonseToResult(json => (json \ "results").validate[Seq[String]])
      }
      .map(_.map(MetricName))
  }

  def listTagNames: Future[Seq[String]] = {
    wsClient
      .url(uri / "api" / "v1" / "tagnames")
      .applyKairosDBAuth
      .get
      .map {
        wsRepsonseToResult(json => (json \ "results").validate[Seq[String]])
      }
  }

  def listTagValues: Future[Seq[String]] = {
    wsClient
      .url(uri / "api" / "v1" / "tagvalues")
      .applyKairosDBAuth
      .get
      .map(
        wsRepsonseToResult(json => (json \ "results").validate[Seq[String]])
      )
  }

  def healthStatus: Future[HealthStatusResults] = {
    wsClient
      .url(uri / "api" / "v1" / "health" / "status")
      .applyKairosDBAuth
      .get
      .map {
        wsRepsonseToResult(_.validate[Seq[String]])
      }
      .map(HealthStatusResults)
  }

  def healthCheck: Future[HealthCheckResult] = {
    wsClient
      .url(uri / "api" / "v1" / "health" / "check")
      .withRequestTimeout(10.seconds)
      .applyKairosDBAuth
      .get
      .map(_.status)
      .map {
        case 204 => AllHealthy
        case _   => Unhealthy
      }
      .recover {
        case _ => Unhealthy
      }
  }

  def version: Future[String] = {
    wsClient
      .url(uri / "api" / "v1" / "version")
      .applyKairosDBAuth
      .get
      .map (
        wsRepsonseToResult(json => (json \ "version").validate[String])
      )
  }

  def deleteMetric(metricName: MetricName): Future[Unit] = {
    wsClient
      .url(uri / "api" / "v1" / "metric" / metricName.name)
      .applyKairosDBAuth
      .delete
      .map(emptyWsRepsonseToResult)
  }

  def addDataPoints(dataPoints: Seq[DataPoint], gzip: Boolean = false): Future[Unit] = {
    val body = Json.toJson(dataPoints)
    logger.debug(Json.prettyPrint(body))

    val (actualBody, contentType) = if(gzip){
      val buff = new ByteArrayOutputStream()
      val gzos = new GZIPOutputStream(buff)
      gzos.write(Json.toBytes(body))
      gzos.finish()
      val gzipBody = buff.toByteArray
      (gzipBody, "application/gzip")
    }else{
      (Json.toBytes(body), "application/json")
    }

    wsClient
      .url(uri / "api" / "v1" / "datapoints")
      .applyKairosDBAuth
      .addHttpHeaders("Content-Type" -> contentType)
      .post(actualBody)
      .map(emptyWsRepsonseToResult)
  }

  def addDataPoint(dataPoint: DataPoint): Future[Unit] = addDataPoints(Seq(dataPoint))

  def queryMetrics(queryMetrics: QueryMetrics): Future[Response] = {
    val query = Json.toJson(queryMetrics)
    logger.debug(Json.prettyPrint(query))
    wsClient
      .url(uri / "api" / "v1" / "datapoints" / "query")
      .applyKairosDBAuth
      .post(query)
      .map {
        wsRepsonseToResult(_.validate[Response])
      }
  }

  /** You can think of this as the exact same as the query but it leaves off the data and just returns the tag information.
    *
    * Note: Currently this is not implemented in the HBase datastore.
    */
  def queryMetricTags(queryMetrics: QueryMetrics): Future[TagsResponse] = {
    val query = Json.toJson(queryMetrics)
    logger.debug(Json.prettyPrint(query))
    wsClient
      .url(uri / "api" / "v1" / "datapoints" / "query" / "tags")
{'fraction_non_alphanumeric': 0.12517939722532292, 'fraction_numerical': 0.0038271408068888536, 'mean_word_length': 3.436858860983375, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, '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': '1458619', 'n_tokens_mistral': 4126, 'n_tokens_neox': 4002, 'n_words': 863}