<commit_before>// Copyright 2016 Foursquare Labs Inc. All Rights Reserved.

package io.fsq.rogue.query

import com.mongodb.{ReadPreference, WriteConcern}
import io.fsq.field.Field
import io.fsq.rogue.{AddLimit, Query, QueryHelpers, QueryOptimizer, Rogue, ShardingOk, !<:<}
import io.fsq.rogue.adapter.MongoClientAdapter
import io.fsq.rogue.types.MongoDisallowed
import scala.collection.generic.CanBuildFrom
import scala.collection.mutable.ArrayBuffer


/** TODO(jacob): All of the collection methods implemented here should get rid of the
  *     option to send down a read preference, and just use the one on the query.
  */
class QueryExecutor[
  MongoCollection[_],
  DocumentValue,
  Document <: java.util.Map[String, DocumentValue],
  MetaRecord,
  Record,
  Result[_]
](
  adapter: MongoClientAdapter[MongoCollection, DocumentValue, Document, MetaRecord, Record, Result],
  optimizer: QueryOptimizer,
  serializer: RogueSerializer[MetaRecord, Record, Document]
) extends Rogue {

  def defaultWriteConcern: WriteConcern = QueryHelpers.config.defaultWriteConcern

  def count[M <: MetaRecord, State](
    query: Query[M, _, State],
    readPreferenceOpt: Option[ReadPreference] = None
  )(
    implicit ev: ShardingOk[M, State],
    ev2: M !<:< MongoDisallowed
  ): Result[Long] = {
    if (optimizer.isEmptyQuery(query)) {
      adapter.wrapEmptyResult(0L)
    } else {
      adapter.count(query, readPreferenceOpt)
    }
  }

  def countDistinct[M <: MetaRecord, FieldType, State](
    query: Query[M, _, State],
    readPreferenceOpt: Option[ReadPreference] = None
  )(
    field: M => Field[FieldType, M]
  )(
    implicit ev: ShardingOk[M, State],
    ev2: M !<:< MongoDisallowed
  ): Result[Long] = {
    if (optimizer.isEmptyQuery(query)) {
      adapter.wrapEmptyResult(0L)
    } else {
      adapter.countDistinct(query, field(query.meta).name, readPreferenceOpt)
    }
  }

  def distinct[M <: MetaRecord, FieldType, State](
    query: Query[M, _, State],
    readPreferenceOpt: Option[ReadPreference] = None
  )(
    field: M => Field[FieldType, M],
    resultTransformer: DocumentValue => FieldType = (_: DocumentValue).asInstanceOf[FieldType]
  )(
    implicit ev: ShardingOk[M, State], ev2: M !<:< MongoDisallowed
  ): Result[Seq[FieldType]] = {
    if (optimizer.isEmptyQuery(query)) {
      adapter.wrapEmptyResult(Vector.empty[FieldType])
    } else {
      adapter.distinct(query, field(query.meta).name, resultTransformer, readPreferenceOpt)
    }
  }

  def insert[R <: Record](
    record: R,
    writeConcern: WriteConcern = defaultWriteConcern
  ): Result[R] = {
    adapter.insert(record, serializer.writeToDocument(record), Some(writeConcern))
  }

  def insertAll[R <: Record](
    records: Seq[R],
    writeConcern: WriteConcern = defaultWriteConcern
  ): Result[Seq[R]] = {
    adapter.insertAll(<commit_msg>new rogue - add missing MongoDisallowed implicits<commit_after>// Copyright 2016 Foursquare Labs Inc. All Rights Reserved.

package io.fsq.rogue.query

import com.mongodb.{ReadPreference, WriteConcern}
import io.fsq.field.Field
import io.fsq.rogue.{AddLimit, Query, QueryHelpers, QueryOptimizer, Rogue, ShardingOk, !<:<}
import io.fsq.rogue.adapter.MongoClientAdapter
import io.fsq.rogue.types.MongoDisallowed
import scala.collection.generic.CanBuildFrom
import scala.collection.mutable.ArrayBuffer


/** TODO(jacob): All of the collection methods implemented here should get rid of the
  *     option to send down a read preference, and just use the one on the query.
  */
class QueryExecutor[
  MongoCollection[_],
  DocumentValue,
  Document <: java.util.Map[String, DocumentValue],
  MetaRecord,
  Record,
  Result[_]
](
  adapter: MongoClientAdapter[MongoCollection, DocumentValue, Document, MetaRecord, Record, Result],
  optimizer: QueryOptimizer,
  serializer: RogueSerializer[MetaRecord, Record, Document]
) extends Rogue {

  def defaultWriteConcern: WriteConcern = QueryHelpers.config.defaultWriteConcern

  def count[M <: MetaRecord, State](
    query: Query[M, _, State],
    readPreferenceOpt: Option[ReadPreference] = None
  )(
    implicit ev: ShardingOk[M, State],
    ev2: M !<:< MongoDisallowed
  ): Result[Long] = {
    if (optimizer.isEmptyQuery(query)) {
      adapter.wrapEmptyResult(0L)
    } else {
      adapter.count(query, readPreferenceOpt)
    }
  }

  def countDistinct[M <: MetaRecord, FieldType, State](
    query: Query[M, _, State],
    readPreferenceOpt: Option[ReadPreference] = None
  )(
    field: M => Field[FieldType, M]
  )(
    implicit ev: ShardingOk[M, State],
    ev2: M !<:< MongoDisallowed
  ): Result[Long] = {
    if (optimizer.isEmptyQuery(query)) {
      adapter.wrapEmptyResult(0L)
    } else {
      adapter.countDistinct(query, field(query.meta).name, readPreferenceOpt)
    }
  }

  def distinct[M <: MetaRecord, FieldType, State](
    query: Query[M, _, State],
    readPreferenceOpt: Option[ReadPreference] = None
  )(
    field: M => Field[FieldType, M],
    resultTransformer: DocumentValue => FieldType = (_: DocumentValue).asInstanceOf[FieldType]
  )(
    implicit ev: ShardingOk[M, State],
    ev2: M !<:< MongoDisallowed
  ): Result[Seq[FieldType]] = {
    if (optimizer.isEmptyQuery(query)) {
      adapter.wrapEmptyResult(Vector.empty[FieldType])
    } else {
      adapter.distinct(query, field(query.meta).name, resultTransformer, readPreferenceOpt)
    }
  }

  def insert[R <: Record](
    record: R,
    writeConcern: WriteConcern = defaultWriteConcern
  )(
    implicit ev: R !<:< MongoDisallowed
  ): Result[R] = {
    adapter.insert(record, serializer.writeToDocument(record), Some(writeConcern))
  }

  def insertAll[R <: Record](
    records: Seq[R],
    writeConcern: WriteConcern = defaultWriteConcern
  )(
    implicit ev: R !<:< MongoDisallowed
  ): Result[Seq[R]] = {
    adapter.insertAll(
{'fraction_non_alphanumeric': 0.12095664143152099, 'fraction_numerical': 0.0030970406056434964, 'mean_word_length': 4.153368794326241, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 9, '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': '3392607', 'n_tokens_mistral': 1929, 'n_tokens_neox': 1866, 'n_words': 410}