<commit_before>package akka.contrib.persistence.mongodb
import akka.NotUsed
import akka.contrib.persistence.mongodb.JournallingFieldNames._
import akka.contrib.persistence.mongodb.RxStreamsInterop._
import akka.persistence.query.{NoOffset, Offset}
import akka.stream.scaladsl._
import akka.stream.stage.{GraphStage, GraphStageLogic, OutHandler}
import akka.stream.{Attributes, Materializer, Outlet, SourceShape}
import com.mongodb.CursorType
import com.mongodb.async.client.Subscription
import org.bson.types.ObjectId
import org.mongodb.scala._
import org.mongodb.scala.bson._
import org.mongodb.scala.model.Aggregates._
import org.mongodb.scala.model.Filters._
import org.mongodb.scala.model.Projections._
import org.mongodb.scala.model.Sorts._

import scala.collection.JavaConverters._
import scala.collection.immutable.Seq
import scala.concurrent.duration._
import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Random, Try}

object CurrentAllEvents {
  def source(driver: ScalaMongoDriver)(implicit m: Materializer): Source[Event, NotUsed] = {
    import driver.ScalaSerializers._
    implicit val ec: ExecutionContext = driver.querySideDispatcher

    Source.fromFuture(driver.journalCollectionsAsFuture)
      .flatMapConcat(_.map(
        _.find()
          .projection(include(EVENTS))
          .asAkka
          .map(e =>
            Option(e.asDocument().get(EVENTS)).filter(_.isArray).map(_.asArray)
              .map(
                _.getValues.asScala.collect{
                  case d:BsonDocument => driver.deserializeJournal(d)
                })
              .getOrElse(Nil)
          ).mapConcat(xs => Seq(xs:_*))
      ).reduceLeftOption(_ concat _)
       .getOrElse(Source.empty))
  }
}

object CurrentPersistenceIds {
  def source(driver: ScalaMongoDriver)(implicit m: Materializer): Source[String, NotUsed] = {
    implicit val ec: ExecutionContext = driver.querySideDispatcher
    val temporaryCollectionName: String = s"persistenceids-${System.currentTimeMillis()}-${Random.nextInt(1000)}"

    Source.fromFuture(for {
      collections <- driver.journalCollectionsAsFuture
      tmpNames    <- Future.sequence(collections.zipWithIndex.map { case (c,idx) =>
                        val nameWithIndex = s"$temporaryCollectionName-$idx"
                        c.aggregate(
                          project(include(PROCESSOR_ID)) ::
                          group(s"$$$PROCESSOR_ID") ::
                          out(nameWithIndex) ::
                          Nil
                        )
                        .asAkka
                        .runWith(Sink.headOption)
                        .map(_ => nameWithIndex)
                      })
      tmps         <- Future.sequence(tmpNames.map(driver.collection))
    } yield tmps )
    .flatMapConcat(_.map(_.find().asAkka).reduceLeftOption(_ ++ _).getOrElse(Source.empty))
    .mapConcat(c => List(c.asDocument().getString("_id").getValue))
    .alsoTo(Sink.onComplete{ _ =>
      driver
        .getCollectionsAsFuture(temporaryCollectionName)
        .foreach(cols =>
          cols.foreach(_.drop().toFuture())
        )
    })
  }
}

object CurrentEventsByPersistenceId {
  def queryFor(persistenceId: String, fromSeq: Long, toSeq: Long): conversions.Bson =
    and(
      equal(PROCESSOR_ID, persistenceId),
      gte(TO, fromSeq),
      lte(FROM, toSeq)
    )

  def source(driver: ScalaMongoDriver, persistenceId: String, fromSeq: Long, toSeq: Long)(implicit m: Materializer): Source[Event, NotUsed] = {
    import driver.ScalaSerializers._
    implicit val ec: ExecutionContext = driver.querySideDispatcher

    val query = queryFor(persistenceId, fromSeq, toSeq)

    Source.fromFuture(driver.getJournal(persistenceId))
      .flatMapConcat(
        _.find(query)
          .sort(ascending(TO))
          .projection(include(EVENTS))
          .asAkka
      ).map(_.asDocument)
       .map(doc =>
        Option(doc.get(EVENTS)).filter(_.isArray).map(_.asArray)
          .map(_.getValues
            .asScala
            .collect{<commit_msg>Source for current Persistence Ids processed in a streamed manner<commit_after>package akka.contrib.persistence.mongodb
import akka.NotUsed
import akka.contrib.persistence.mongodb.JournallingFieldNames._
import akka.contrib.persistence.mongodb.RxStreamsInterop._
import akka.persistence.query.{NoOffset, Offset}
import akka.stream.scaladsl._
import akka.stream.stage.{GraphStage, GraphStageLogic, OutHandler}
import akka.stream.{Attributes, Materializer, Outlet, SourceShape}
import com.mongodb.CursorType
import com.mongodb.async.client.Subscription
import org.bson.types.ObjectId
import org.mongodb.scala._
import org.mongodb.scala.bson._
import org.mongodb.scala.model.Aggregates._
import org.mongodb.scala.model.Filters._
import org.mongodb.scala.model.Projections._
import org.mongodb.scala.model.Sorts._

import scala.collection.JavaConverters._
import scala.collection.immutable.Seq
import scala.concurrent.ExecutionContext
import scala.concurrent.duration._
import scala.util.Try

object CurrentAllEvents {
  def source(driver: ScalaMongoDriver)(implicit m: Materializer): Source[Event, NotUsed] = {
    import driver.ScalaSerializers._
    implicit val ec: ExecutionContext = driver.querySideDispatcher

    Source.fromFuture(driver.journalCollectionsAsFuture)
      .flatMapConcat(_.map(
        _.find()
          .projection(include(EVENTS))
          .asAkka
          .map(e =>
            Option(e.asDocument().get(EVENTS)).filter(_.isArray).map(_.asArray)
              .map(
                _.getValues.asScala.collect{
                  case d:BsonDocument => driver.deserializeJournal(d)
                })
              .getOrElse(Nil)
          ).mapConcat(xs => Seq(xs:_*))
      ).reduceLeftOption(_ concat _)
       .getOrElse(Source.empty))
  }
}

object CurrentPersistenceIds {
  def source(driver: ScalaMongoDriver)(implicit m: Materializer): Source[String, NotUsed] = {
    implicit val ec: ExecutionContext = driver.querySideDispatcher

    Source.fromFuture(driver.journalCollectionsAsFuture)
      .mapConcat(identity)
      .flatMapConcat(_.aggregate(project(include(PROCESSOR_ID)) :: group(s"$$$PROCESSOR_ID") :: Nil).asAkka)
      .map { doc => Option(doc.getString("_id").getValue).getOrElse("") }
      .filterNot(_.isEmpty)
  }
}

object CurrentEventsByPersistenceId {
  def queryFor(persistenceId: String, fromSeq: Long, toSeq: Long): conversions.Bson =
    and(
      equal(PROCESSOR_ID, persistenceId),
      gte(TO, fromSeq),
      lte(FROM, toSeq)
    )

  def source(driver: ScalaMongoDriver, persistenceId: String, fromSeq: Long, toSeq: Long)(implicit m: Materializer): Source[Event, NotUsed] = {
    import driver.ScalaSerializers._
    implicit val ec: ExecutionContext = driver.querySideDispatcher

    val query = queryFor(persistenceId, fromSeq, toSeq)

    Source.fromFuture(driver.getJournal(persistenceId))
      .flatMapConcat(
        _.find(query)
          .sort(ascending(TO))
          .projection(include(EVENTS))
          .asAkka
      ).map(_.asDocument)
       .map(doc =>
        Option(doc.get(EVENTS)).filter(_.isArray).map(_.asArray)
          .map(_.getValues
            .asScala
            .collect{
{'fraction_non_alphanumeric': 0.1118648760902672, 'fraction_numerical': 0.0005537865152983525, 'mean_word_length': 3.0066555740432612, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 17, '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': '5111998', 'n_tokens_mistral': 2219, 'n_tokens_neox': 2165, 'n_words': 359}