Cassandra closeAsync

locationtech · Jul 10, 2022 · a4372ef · a4372ef
1 parent 5141587
commit a4372ef
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 9 deletions.
diff --git a/cassandra-spark/src/main/scala/geotrellis/spark/store/cassandra/CassandraRDDReader.scala b/cassandra-spark/src/main/scala/geotrellis/spark/store/cassandra/CassandraRDDReader.scala
@@ -89,7 +89,7 @@ object CassandraRDDReader {
           }
 
           /** Close partition session */
-          (result ++ Iterator({ session.close(); Seq.empty[(K, V)] })).flatten
+          (result ++ Iterator({ session.closeAsync(); Seq.empty[(K, V)] })).flatten
         }
       }
   }

diff --git a/cassandra-spark/src/main/scala/geotrellis/spark/store/cassandra/CassandraRDDWriter.scala b/cassandra-spark/src/main/scala/geotrellis/spark/store/cassandra/CassandraRDDWriter.scala
@@ -151,7 +151,7 @@ object CassandraRDDWriter {
                 .flatMap(rowToBytes)
                 .map(retire)
                 .parJoinUnbounded
-                .onComplete { fs2.Stream eval session.closeF[IO] }
+                .onComplete { fs2.Stream eval IO(session.closeAsync) }
 
               results
                 .compile

diff --git a/cassandra-spark/src/test/resources/application.conf b/cassandra-spark/src/test/resources/application.conf
@@ -4,4 +4,4 @@ datastax-java-driver {
       local-datacenter = datacenter1
     }
   }
-}
+}
diff --git a/cassandra/src/main/scala/geotrellis/store/cassandra/CassandraInstance.scala b/cassandra/src/main/scala/geotrellis/store/cassandra/CassandraInstance.scala
@@ -77,12 +77,14 @@ trait CassandraInstance extends Serializable {
   /** With session close */
   def withSessionDo[T](block: CqlSession => T): T = {
     val session = getSession()
-    try block(session) finally session.close()
+    try block(session) finally session.closeAsync()
   }
 
-  def closeAsync[F[_]: Async]: F[Unit] = session.closeF
+  // def closeAsync[F[_]: Async]: F[Unit] = session.closeF
 
-  def close(): Unit = session.close()
+  // def close(): Unit = session.close()
+
+  def closeAsync = session.closeAsync()
 }
 
 case class BaseCassandraInstance(
@@ -134,7 +136,7 @@ object Cassandra {
   implicit def instanceToSession[T <: CassandraInstance](instance: T): CqlSession = instance.session
 
   def withCassandraInstance[T <: CassandraInstance, K](instance: T)(block: T => K): K = block(instance)
-  def withCassandraInstanceDo[T <: CassandraInstance, K](instance: T)(block: T => K): K = try block(instance) finally instance.close()
+  def withCassandraInstanceDo[T <: CassandraInstance, K](instance: T)(block: T => K): K = try block(instance) finally instance.closeAsync()
 
   def withBaseCassandraInstance[K](hosts: Seq[String],
                                    username: String,
@@ -155,7 +157,7 @@ object Cassandra {
                                      password: String,
                                      cassandraConfig: CassandraConfig)(block: CassandraInstance => K): K = {
     val instance = BaseCassandraInstance(hosts, username, password, cassandraConfig)
-    try block(instance) finally instance.close()
+    try block(instance) finally instance.closeAsync()
   }
   def withBaseCassandraInstanceDo[K](hosts: Seq[String],
                                      username: String,

diff --git a/cassandra/src/main/scala/geotrellis/store/cassandra/package.scala b/cassandra/src/main/scala/geotrellis/store/cassandra/package.scala
@@ -24,7 +24,7 @@ import com.datastax.oss.driver.api.core.cql.{AsyncResultSet, Statement}
 import java.math.BigInteger
 import java.util.concurrent.CompletableFuture
 
-package object cassandra {
+package object cassandra extends Serializable {
   implicit def bigToBig(i: BigInt): BigInteger = new BigInteger(i.toByteArray)
 
   implicit class CompletableFutureOps[T](val self: CompletableFuture[T]) extends AnyVal {
@@ -35,4 +35,9 @@ package object cassandra {
     def closeF[F[_]: Async]: F[Unit] = self.closeAsync().toCompletableFuture.liftTo.void
     def executeF[F[_]: Async](statement: Statement[_]): F[AsyncResultSet] = self.executeAsync(statement).toCompletableFuture.liftTo
   }
+
+  implicit class AsyncResultSetOps(val self: AsyncResultSet) extends AnyVal {
+    def nonEmpty: Boolean = self.remaining() > 0 || self.hasMorePages
+    def isEmpty: Boolean = !nonEmpty
+  }
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -89,7 +89,7 @@ object CassandraRDDReader { @@
               }
               /** Close partition session */
-              (result ++ Iterator({ session.close(); Seq.empty[(K, V)] })).flatten
+              (result ++ Iterator({ session.closeAsync(); Seq.empty[(K, V)] })).flatten
             }
           }
       }
@@ Expand Down @@