Cassandra sink

cassandra/src/main/scala/akka/stream/alpakka/cassandra/CassandraSourceStage.scala

@@ -6,11 +6,12 @@ package akka.stream.alpakka.cassandra
 import akka.stream._
 import akka.stream.stage.{ AsyncCallback, GraphStage, GraphStageLogic, OutHandler }
 import com.datastax.driver.core.{ ResultSet, Row, Session, Statement }
-import com.google.common.util.concurrent.{ FutureCallback, Futures, ListenableFuture }
 
-import scala.concurrent.{ Future, Promise }
+import scala.concurrent.Future
 import scala.util.{ Failure, Success, Try }
 
+import akka.stream.alpakka.cassandra.cassandra._
+
 class CassandraSourceStage(futStmt: Future[Statement], session: Session) extends GraphStage[SourceShape[Row]] {
   val out: Outlet[Row] = Outlet("CassandraSource.out")
   override val shape: SourceShape[Row] = SourceShape(out)
@@ -25,7 +26,7 @@ class CassandraSourceStage(futStmt: Future[Statement], session: Session) extends
 
         futFetchedCallback = getAsyncCallback[Try[ResultSet]](tryPushAfterFetch)
 
-        val futRs = futStmt.flatMap(stmt => guavaFutToScalaFut(session.executeAsync(stmt)))
+        val futRs = futStmt.flatMap(stmt => session.executeAsync(stmt).asScala())
         futRs.onComplete(futFetchedCallback.invoke)
       }
 
@@ -38,7 +39,7 @@ class CassandraSourceStage(futStmt: Future[Statement], session: Session) extends
             case Some(rs) if rs.isExhausted                     => completeStage()
             case Some(rs) =>
               // fetch next page
-              val futRs = guavaFutToScalaFut(rs.fetchMoreResults())
+              val futRs = rs.fetchMoreResults().asScala()
               futRs.onComplete(futFetchedCallback.invoke)
             case None => () // doing nothing, waiting for futRs in preStart() to be completed
           }
@@ -59,15 +60,5 @@ class CassandraSourceStage(futStmt: Future[Statement], session: Session) extends
         case Failure(failure) => failStage(failure)
       }
     }
-
-  private def guavaFutToScalaFut[A](guavaFut: ListenableFuture[A]): Future[A] = {
-    val p = Promise[A]()
-    val callback = new FutureCallback[A] {
-      override def onSuccess(a: A): Unit = p.success(a)
-      override def onFailure(err: Throwable): Unit = p.failure(err)
-    }
-    Futures.addCallback(guavaFut, callback)
-    p.future
-  }
 }
 

cassandra/src/main/scala/akka/stream/alpakka/cassandra/package.scala

@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2016 Lightbend Inc. <http://www.lightbend.com>
+ */
+package akka.stream.alpakka.cassandra
+
+import com.google.common.util.concurrent.{ FutureCallback, Futures, ListenableFuture }
+
+import scala.concurrent.{ Future, Promise }
+
+package object cassandra {
+  implicit class GuavaFutureOpts[A](val guavaFut: ListenableFuture[A]) extends AnyVal {
+    def asScala(): Future[A] = {
+      val p = Promise[A]()
+      val callback = new FutureCallback[A] {
+        override def onSuccess(a: A): Unit = p.success(a)
+        override def onFailure(err: Throwable): Unit = p.failure(err)
+      }
+      Futures.addCallback(guavaFut, callback)
+      p.future
+    }
+  }
+}

cassandra/src/main/scala/akka/stream/alpakka/cassandra/scaladsl/CassandraSink.scala

@@ -0,0 +1,18 @@
+/*
+ * Copyright (C) 2016 Lightbend Inc. <http://www.lightbend.com>
+ */
+package akka.stream.alpakka.cassandra.scaladsl
+
+import akka.Done
+import akka.stream.scaladsl.{ Flow, Keep, Sink }
+import com.datastax.driver.core.{ BoundStatement, PreparedStatement, Session }
+
+import scala.concurrent.{ ExecutionContext, Future }
+
+import akka.stream.alpakka.cassandra.cassandra._
+
+object CassandraSink {
+  def apply[T](parallelism: Int, statement: PreparedStatement, statementBinder: (T, PreparedStatement) => BoundStatement)(implicit session: Session, ex: ExecutionContext): Sink[T, Future[Done]] =
+    Flow[T].mapAsyncUnordered(parallelism)(t ⇒
+      session.executeAsync(statementBinder(t, statement)).asScala()).toMat(Sink.ignore)(Keep.right)
+}

cassandra/src/main/scala/akka/stream/alpakka/cassandra/scaladsl/CassandraSource.scala

@@ -3,10 +3,10 @@
  */
 package akka.stream.alpakka.cassandra.scaladsl
 
-import akka.NotUsed
+import akka.{ Done, NotUsed }
 import akka.stream.alpakka.cassandra.CassandraSourceStage
-import akka.stream.scaladsl.Source
-import com.datastax.driver.core.{ Row, Session, Statement }
+import akka.stream.scaladsl.{ Sink, Source }
+import com.datastax.driver.core._
 
 import scala.concurrent.Future
 

cassandra/src/test/scala/akka/stream/alpakka/cassandra/scaladsl/CassandraSourceSpec.scala

@@ -5,14 +5,16 @@ package akka.stream.alpakka.cassandra.scaladsl
 
 import akka.actor.ActorSystem
 import akka.stream.ActorMaterializer
-import akka.stream.scaladsl.Sink
+import akka.stream.scaladsl.{ Sink, Source }
 import com.datastax.driver.core.{ Cluster, SimpleStatement }
 import org.scalatest._
 import org.scalatest.concurrent.ScalaFutures
 
 import scala.concurrent._
 import scala.concurrent.duration._
 
+import scala.collection.JavaConverters._
+
 /**
  * All the tests must be run with a local Cassandra running on default port 9042.
  */
@@ -114,6 +116,19 @@ class CassandraSourceSpec extends WordSpec with ScalaFutures with BeforeAndAfter
       rows mustBe empty
     }
 
-  }
+    "sink should write to the table" in {
+      import system.dispatcher
+
+      val source = Source(0 to 10)
+      val sink = CassandraSink[Integer](parallelism = 2, session.prepare("INSERT INTO akka_stream_scala_test.test(id) VALUES (?)"), (t, p) => p.bind(t))
+
+      val result = source.map(i => i: Integer).runWith(sink)
 
+      result.futureValue
+
+      val found = session.execute("select id from akka_stream_scala_test.test").all().asScala.map(_.getInt("id"))
+
+      found.toSet mustBe (0 to 10).toSet
+    }
+  }
 }