From d449c66ed70b1e6b65d424cb7c0f255743b1cd2a Mon Sep 17 00:00:00 2001 From: "nikita.myazin" Date: Thu, 8 Dec 2022 13:58:19 +0200 Subject: [PATCH 1/2] add Session#execute for Batch --- .../zio/cassandra/session/SessionSpec.scala | 19 +++++++++++++++++-- .../scala/zio/cassandra/session/Session.scala | 5 ++++- .../cassandra/session/cql/query/Batch.scala | 7 +++++-- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/it/scala/zio/cassandra/session/SessionSpec.scala b/src/it/scala/zio/cassandra/session/SessionSpec.scala index 8defd0a..9a082d4 100644 --- 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 index 67eaf7a..dcb6055 100644 --- 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 index 12b5662..4eda4bd 100644 --- 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,18 @@ 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 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 = From 43bde9cf33c3308b5f2702d0eadfa767ad5297ed Mon Sep 17 00:00:00 2001 From: "nikita.myazin" Date: Thu, 8 Dec 2022 14:18:24 +0200 Subject: [PATCH 2/2] embrace Batch mutability --- src/main/scala/zio/cassandra/session/cql/query/Batch.scala | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/main/scala/zio/cassandra/session/cql/query/Batch.scala b/src/main/scala/zio/cassandra/session/cql/query/Batch.scala index 4eda4bd..4d1289b 100644 --- a/src/main/scala/zio/cassandra/session/cql/query/Batch.scala +++ b/src/main/scala/zio/cassandra/session/cql/query/Batch.scala @@ -6,7 +6,10 @@ 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()