Merge pull request #30 from narma/add_session_batch_execute

Add session batch execute
narma · Dec 8, 2022 · 3f82de1 · 3f82de1
2 parents 8f74316 + b0d9521
commit 3f82de1
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 8 deletions.
diff --git a/src/it/scala/zio/cassandra/session/SessionSpec.scala b/src/it/scala/zio/cassandra/session/SessionSpec.scala
@@ -5,6 +5,7 @@ import com.datastax.oss.driver.api.core.cql.SimpleStatement
 import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException
 import com.dimafeng.testcontainers.CassandraContainer
 import zio.cassandra.session.cql.CqlStringContext
+import zio.cassandra.session.cql.query.Batch
 import zio.cassandra.session.cql.unsafe.lift
 import zio.test.Assertion._
 import zio.test._
@@ -104,7 +105,7 @@ object SessionSpec extends ZIOCassandraSpec with ZIOCassandraSpecUtils {
     test("selectFirst should return Some(null) for null") {
       for {
         session <- ZIO.service[Session]
-        result <- session.selectFirst(s"select data FROM $keyspace.test_data WHERE id = 0").map(_.map(_.getString(0)))
+        result  <- session.selectFirst(s"select data FROM $keyspace.test_data WHERE id = 0").map(_.map(_.getString(0)))
       } yield assertTrue(result.contains(null))
     },
     test("selectFirst interpolated query (cqlConst) should return Some") {
@@ -138,9 +139,23 @@ object SessionSpec extends ZIOCassandraSpec with ZIOCassandraSpecUtils {
     test("execute will create a table") {
       for {
         session <- ZIO.service[Session]
-        table   = "table_" + UUID.randomUUID().toString.replaceAll("-", "_")
+        table    = "table_" + UUID.randomUUID().toString.replaceAll("-", "_")
         created <- session.execute(cqlConst"create table $keyspace.$table(id text primary key)")
       } yield assertTrue(created)
+    },
+    test("execute will insert batched data") {
+      for {
+        session  <- ZIO.service[Session]
+        tbl       = "table_" + UUID.randomUUID().toString.replaceAll("-", "_")
+        table     = s"$keyspace.$tbl"
+        _        <- session.execute(cqlConst"create table $table(id text primary key)")
+        insert1  <- session.prepare(cqlConst"insert into $table(id) values ('primary key 1')")
+        insert2  <- session.prepare(cqlConst"insert into $table(id) values ('primary key 2')")
+        insert3  <- session.prepare(cqlConst"insert into $table(id) values ('primary key 3')")
+        batch     = Batch.unlogged.add(Seq(insert1, insert2, insert3))
+        inserted <- session.execute(batch)
+        result   <- session.selectFirst(cqlConst"select count(*) from $table".as[Long])
+      } yield assertTrue(inserted, result.contains(3))
     }
   )
 }
diff --git a/src/main/scala/zio/cassandra/session/Session.scala b/src/main/scala/zio/cassandra/session/Session.scala
@@ -6,7 +6,7 @@ import com.datastax.oss.driver.api.core.metadata.Metadata
 import com.datastax.oss.driver.api.core.metrics.Metrics
 import com.datastax.oss.driver.api.core.{ CqlIdentifier, CqlSession, CqlSessionBuilder }
 import zio._
-import zio.cassandra.session.cql.query.{ PreparedQuery, QueryTemplate }
+import zio.cassandra.session.cql.query.{ Batch, PreparedQuery, QueryTemplate }
 import zio.stream.ZStream.Pull
 import zio.stream.{ Stream, ZStream }
 
@@ -34,6 +34,9 @@ trait Session {
   final def execute(template: QueryTemplate[_]): Task[Boolean] =
     prepare(template).flatMap(_.execute)
 
+  final def execute(batch: Batch): Task[Boolean] =
+    execute(batch.build).map(_.wasApplied)
+
   final def select[R](template: QueryTemplate[R]): Stream[Throwable, R] =
     ZStream.fromZIO(prepare(template)).flatMap(_.select)
 

diff --git a/src/main/scala/zio/cassandra/session/cql/query/Batch.scala b/src/main/scala/zio/cassandra/session/cql/query/Batch.scala
@@ -1,16 +1,20 @@
 package zio.cassandra.session.cql.query
 
-import com.datastax.oss.driver.api.core.cql.{ BatchStatementBuilder, BatchType }
+import com.datastax.oss.driver.api.core.cql.{ BatchStatement, BatchStatementBuilder, BatchType }
 import zio.cassandra.session.Session
 import zio.{ RIO, ZIO }
 
 class Batch(batchStatementBuilder: BatchStatementBuilder) {
-  def add(queries: Seq[PreparedQuery[_]]) = new Batch(batchStatementBuilder.addStatements(queries.map(_.statement): _*))
+
+  def add(queries: Seq[PreparedQuery[_]]): this.type = {
+    batchStatementBuilder.addStatements(queries.map(_.statement): _*)
+    this
+  }
+
+  def build: BatchStatement = batchStatementBuilder.build()
 
   def execute: RIO[Session, Boolean] =
-    ZIO.serviceWithZIO { session =>
-      session.execute(batchStatementBuilder.build()).map(_.wasApplied)
-    }
+    ZIO.serviceWithZIO(_.execute(this))
 
   def config(config: BatchStatementBuilder => BatchStatementBuilder): Batch =
     new Batch(config(batchStatementBuilder))