Merge pull request #29 from narma/add_session_batch_execute_zio-1.x

Add session batch execute zio 1.x
narma · Dec 8, 2022 · 49054cf · 49054cf
2 parents 766cff6 + 43bde9c
commit 49054cf
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 6 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 CassandraSpecUtils {
     testM("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))
     },
     testM("selectFirst interpolated query (cqlConst) should return Some") {
@@ -131,9 +132,23 @@ object SessionSpec extends CassandraSpecUtils {
     testM("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)
+    },
+    testM("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.Stream
 import zio.stream.ZStream.Pull
 
@@ -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] =
     Stream.fromEffect(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,15 +1,21 @@
 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.{ Has, 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[Has[Session], Boolean] =
     ZIO.accessM[Has[Session]] { session =>
-      session.get.execute(batchStatementBuilder.build()).map(_.wasApplied)
+      session.get.execute(this)
     }
 
   def config(config: BatchStatementBuilder => BatchStatementBuilder): Batch =