<commit_before>package com.mnubo
package dbschemas

import java.text.SimpleDateFormat
import java.util.Date

import com.datastax.driver.core.{Session, Cluster}
import com.mnubo.test_utils.cassandra.DockerCassandra
import com.typesafe.config.Config
import org.joda.time.{DateTimeZone, DateTime}

import scala.collection.JavaConverters._
import scala.util.control.NonFatal

object CassandraDatabase extends Database {
  val name = "cassandra"

  override def openConnection(schemaName: String,
                              hosts: String,
                              port: Int,
                              userName: String,
                              pwd: String,
                              createDatabaseStatement: String,
                              config: Config): DatabaseConnection =
    new CassandraConnection(
      schemaName,
      hosts,
      if (port > 0) port else 9042,
      createDatabaseStatement,
      config.getInt("max_schema_agreement_wait_seconds"),
      config.getBoolean("force_pull_verification_db")
    )

  override def testDockerBaseImage =
    DatabaseDockerImage(
      name        = "spotify/cassandra",
      exposedPort = 9042,
      isStarted   = (log, _) => log.contains("Listening for thrift clients...")
    )
}

class CassandraConnection(
                           schemaName: String,
                           hosts: String,
                           port: Int,
                           createDatabaseStatement: String,
                           maxSchemaAgreementWaitSeconds: Int,
                           forcePullVerificationDb: Boolean) extends DatabaseConnection {
  private val cluster = Cluster
    .builder()
    .addContactPoints(hosts.split(","): _*)
    .withMaxSchemaAgreementWaitSeconds(maxSchemaAgreementWaitSeconds)
    .withPort(port)
    .build()
  private val session = cluster.connect()
  private val df = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSZ")
  private var keyspace: String = null

  override def setActiveSchema(keyspace: String) {
    this.keyspace = keyspace
    if (!hasKeyspace) execute(createDatabaseStatement.replace("@@DATABASE_NAME@@", keyspace))
    execute("USE " + keyspace)
  }

  override def execute(smt: String): Unit =
    session.execute(smt)

  override def innerConnection: AnyRef =
    session

  /** For tests, or QA, we might want to recreate a database instance from scratch. Implementors should know how to properly clean an existing database. */
  override def dropDatabase() = {
    cluster<commit_msg>Set Cassandra consistency level to ALL, in order to avoid race conditions during schema creation.<commit_after>package com.mnubo
package dbschemas

import java.text.SimpleDateFormat
import java.util.Date

import com.datastax.driver.core.{ConsistencyLevel, SimpleStatement, Session, Cluster}
import com.mnubo.test_utils.cassandra.DockerCassandra
import com.typesafe.config.Config
import org.joda.time.{DateTimeZone, DateTime}

import scala.collection.JavaConverters._
import scala.util.control.NonFatal

object CassandraDatabase extends Database {
  val name = "cassandra"

  override def openConnection(schemaName: String,
                              hosts: String,
                              port: Int,
                              userName: String,
                              pwd: String,
                              createDatabaseStatement: String,
                              config: Config): DatabaseConnection =
    new CassandraConnection(
      schemaName,
      hosts,
      if (port > 0) port else 9042,
      createDatabaseStatement,
      config.getInt("max_schema_agreement_wait_seconds"),
      config.getBoolean("force_pull_verification_db")
    )

  override def testDockerBaseImage =
    DatabaseDockerImage(
      name        = "spotify/cassandra",
      exposedPort = 9042,
      isStarted   = (log, _) => log.contains("Listening for thrift clients...")
    )
}

class CassandraConnection(
                           schemaName: String,
                           hosts: String,
                           port: Int,
                           createDatabaseStatement: String,
                           maxSchemaAgreementWaitSeconds: Int,
                           forcePullVerificationDb: Boolean) extends DatabaseConnection {
  private val cluster = Cluster
    .builder()
    .addContactPoints(hosts.split(","): _*)
    .withMaxSchemaAgreementWaitSeconds(maxSchemaAgreementWaitSeconds)
    .withPort(port)
    .build()
  private val session = cluster.connect()
  private val df = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSZ")
  private var keyspace: String = null

  override def setActiveSchema(keyspace: String) {
    this.keyspace = keyspace
    if (!hasKeyspace) execute(createDatabaseStatement.replace("@@DATABASE_NAME@@", keyspace))
    execute("USE " + keyspace)
  }

  override def execute(smt: String): Unit =
    session.execute(new SimpleStatement(smt).setConsistencyLevel(ConsistencyLevel.ALL))

  override def innerConnection: AnyRef =
    session

  /** For tests, or QA, we might want to recreate a database instance from scratch. Implementors should know how to properly clean an existing database. */
  override def dropDatabase() = {
    cluster
{'fraction_non_alphanumeric': 0.08071920428462127, 'fraction_numerical': 0.0034429992348890587, 'mean_word_length': 2.4698075646980757, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '6815620', 'n_tokens_mistral': 1408, 'n_tokens_neox': 1355, 'n_words': 342}