Managed stream & advanced aggregation

Author: cchantep

documentation/manual/scalaGuide/main/sql/ScalaAnorm.md

@@ -206,30 +206,26 @@ val code: String = SQL"""
 
 This feature tries to make faster, more concise and easier to read the way to retrieve data in Anorm. Please, feel free to use it wherever you see a combination of `SQL().on()` functions (or even an only `SQL()` without parameters).
 
-## Retrieving data using the Stream API
+## Streaming results
 
-The first way to access the results of a select query is to use the Stream API.
+Query results can be processed as stream of row, not having all loaded in memory.
 
-When you call `apply()` on any SQL statement, you will receive a lazy `Stream` of `Row` instances, where each row can be seen as a dictionary:
+In the following example we will count the number of country rows.
 
 ```scala
-// Create an SQL query
-val selectCountries = SQL("Select * from Country")
- 
-// Transform the resulting Stream[Row] to a List[(String,String)]
-val countries = selectCountries().map(row => 
-  row[String]("code") -> row[String]("name")
-).toList
+val countryCount: Long = 
+  SQL"Select count(*) as c from Country".fold(0l) { (c, _) => c + 1 }
 ```
 
-In the following example we will count the number of `Country` entries in the database, so the result set will be a single row with a single column:
+It's also possible to partially treat the row stream.
 
 ```scala
-// First retrieve the first row
-val firstRow = SQL("Select count(*) as c from Country").apply().head
- 
-// Next get the content of the 'c' column as Long
-val countryCount = firstRow[Long]("c")
+// Create an SQL query
+val books: List[String] = SQL("Select name from Books").
+  foldWhile(List[String]()) { (list, row) => 
+    if (list.size == 100) (list -> false) // stop with `list`
+    else (list := row[String]("name")) -> true // continue with one more name
+  }
 ```
 
 ### Multi-value support

framework/src/anorm/src/main/scala/anorm/Anorm.scala

@@ -264,9 +264,32 @@ private[anorm] trait Sql {
   /**
    * Executes this SQL statement as query, returns result as Row stream.
    */
+  @deprecated(message =
+    "Use [[fold]] instead, which manages resources and memory", "2.4")
   def apply()(implicit connection: Connection): Stream[Row] =
     Sql.resultSetToStream(resultSet())
 
+  /**
+   * Aggregates over the whole stream of row using the specified operator.
+   *
+   * @param z the start value
+   * @param op Aggregate operator
+   * @return Either list of failures at left, or aggregated value
+   * @see #foldWhile
+   */
+  def fold[T](z: T)(op: (T, Row) => T)(implicit connection: Connection): Either[List[Throwable], T] =
+    Sql.fold(resultSet())(z) { (t, r) => op(t, r) -> true } acquireFor identity
+
+  /**
+   * Aggregates over part of or the while row stream, using the specified operator.
+   *
+   * @param z the start value
+   * @param op Aggregate operator. Returns aggregated value along with true if aggregation must process next value, or false to stop with current value.
+   * @return Either list of failures at left, or aggregated value
+   */
+  def foldWhile[T](z: T)(op: (T, Row) => (T, Boolean))(implicit connection: Connection): Either[List[Throwable], T] =
+    Sql.fold(resultSet())(z) { (t, r) => op(t, r) } acquireFor identity
+
   /**
    * Executes this statement as query (see [[executeQuery]]) and returns result.
    */
@@ -364,24 +387,31 @@ object Sql { // TODO: Rename to SQL
   }
 
   private[anorm] def as[T](parser: ResultSetParser[T], rs: ManagedResource[ResultSet])(implicit connection: Connection): T =
-    parser(Sql.resultSetToStream(rs)) match {
+    parser(Sql.resultSetToStream(rs) /* TODO: Refactory with fold */ ) match {
       case Success(a) => a
       case Error(e) => sys.error(e.toString)
     }
 
-  private def fold[T](res: ManagedResource[ResultSet])(initial: T)(f: (T, Row) => T): ManagedResource[T] = res map { rs =>
+  /**
+   * @param f Aggregate operator. Returns aggregated value along with true if aggregation must process next value, or false to stop with current value.
+   */
+  private[anorm] def fold[T](res: ManagedResource[ResultSet])(initial: T)(f: (T, Row) => (T, Boolean)): ManagedResource[T] = res map { rs =>
     val rsMetaData: MetaData = metaData(rs)
     val columns: List[Int] = List.range(1, rsMetaData.columnCount + 1)
     @inline def data(rs: ResultSet) = columns.map(rs.getObject(_))
 
     @annotation.tailrec
     def go(r: ResultSet, s: T): T =
-      if (r.next()) go(r, f(s, new SqlRow(rsMetaData, data(rs)))) else s
+      if (r.next()) {
+        val (v, cont) = f(s, new SqlRow(rsMetaData, data(rs)))
+        if (cont) go(r, v) else v
+      } else s
 
     go(rs, initial)
   }
 
-  private[anorm] def resultSetToStream(rs: ManagedResource[ResultSet]): Stream[Row] = fold(rs)(Stream.empty[Row])((s, r) => s :+ r) acquireAndGet identity
+  // TODO: Remove when unused
+  private[anorm] def resultSetToStream(rs: ManagedResource[ResultSet]): Stream[Row] = fold(rs)(Stream.empty[Row])((s, r) => (s :+ r) -> true) acquireAndGet identity
 
   private case class SqlRow(metaData: MetaData, data: List[Any]) extends Row {
     override lazy val toString = "Row(" + metaData.ms.zip(data).map(t => s"'${t._1.column}': ${t._2} as ${t._1.clazz}").mkString(", ") + ")"

framework/src/anorm/src/main/scala/anorm/SqlQueryResult.scala

@@ -40,17 +40,43 @@ final case class SqlQueryResult(
       _.headOption.map(new SQLWarning(_)), Option(_))
 
   /** Returns stream of row from query result. */
+  @deprecated("Use [[fold]] or [[foldWhile]] instead, which manages resources and memory", "2.4")
   def apply()(implicit connection: Connection): Stream[Row] =
     Sql.resultSetToStream(resultSet)
 
+  /**
+   * Aggregates over the whole row stream using the specified operator.
+   *
+   * @param z the start value
+   * @param op Aggregate operator
+   * @return Either list of failures at left, or aggregated value
+   * @see #foldWhile
+   */
+  def fold[T](z: T)(op: (T, Row) => T)(implicit connection: Connection): Either[List[Throwable], T] =
+    Sql.fold(resultSet)(z) { (t, r) => op(t, r) -> true } acquireFor identity
+
+  /**
+   * Aggregates over part of or the while row stream,
+   * using the specified operator.
+   *
+   * @param z the start value
+   * @param op Aggregate operator. Returns aggregated value along with true if aggregation must process next value, or false to stop with current value.
+   * @return Either list of failures at left, or aggregated value
+   */
+  def foldWhile[T](z: T)(op: (T, Row) => (T, Boolean))(implicit connection: Connection): Either[List[Throwable], T] =
+    Sql.fold(resultSet)(z) { (t, r) => op(t, r) } acquireFor identity
+
   def as[T](parser: ResultSetParser[T])(implicit connection: Connection): T =
     Sql.as(parser, resultSet)
 
+  // TODO: Scaladoc as `as` equivalent
   def list[A](rowParser: RowParser[A])(implicit connection: Connection): Seq[A] = as(rowParser.*)
 
+  // TODO: Scaladoc as `as` equivalent
   def single[A](rowParser: RowParser[A])(implicit connection: Connection): A =
     as(ResultSetParser.single(rowParser))
 
+  // TODO: Scaladoc as `as` equivalent
   def singleOpt[A](rowParser: RowParser[A])(implicit connection: Connection): Option[A] = as(ResultSetParser.singleOpt(rowParser))
 
   @deprecated(message = "Use [[as]]", since = "2.3.2")

framework/src/anorm/src/test/scala/anorm/AnormSpec.scala

@@ -321,6 +321,76 @@ object AnormSpec extends Specification with H2Database with AnormTest {
       }
   }
 
+  "Aggregation over all rows" should {
+    "be empty when there is no result" in withQueryResult(QueryResult.Nil) {
+      implicit c => 
+      SQL"EXEC test".fold[Option[Int]](None)({ (_, _) => Some(0) }).
+        aka("aggregated value") must beRight(None)
+
+    }
+
+    "be parsed from mapped result" in withQueryResult(
+      rowList2(classOf[String] -> "foo", classOf[Int] -> "bar").
+        append("row1", 100) :+ ("row2", 200)) { implicit c =>
+
+        SQL"SELECT * FROM test".fold(List[(String,Int)]())(
+          { (l, row) => l :+ (row[String]("foo") -> row[Int]("bar")) }).
+          aka("tuple stream") must_== Right(List("row1" -> 100, "row2" -> 200))
+
+      }
+
+    "handle failure" in withQueryResult(
+      rowList1(classOf[String] -> "foo") :+ "A" :+ "B") { implicit c =>
+      var i = 0
+        SQL"SELECT str".fold(Set[String]()) { (l, row) => 
+      if (i == 0) { i = i+1; l + row[String]("foo") } else sys.error("Failure")
+
+      } aka "aggregate on failure" must beLike {
+        case Left(err :: Nil) => err.getMessage aka "failure" must_== "Failure"
+      } and (i aka "row count" must_== 1)
+    }
+  }
+
+  "Aggregation over variable number of rows" should {
+    "be empty when there is no result" in withQueryResult(QueryResult.Nil) {
+      implicit c => 
+      SQL"EXEC test".foldWhile[Option[Int]](None)(
+        { (_, _) => Some(0) -> true }).
+        aka("aggregated value") must beRight(None)
+
+    }
+
+    "be parsed from mapped result" in withQueryResult(
+      rowList2(classOf[String] -> "foo", classOf[Int] -> "bar").
+        append("row1", 100) :+ ("row2", 200)) { implicit c =>
+
+        SQL"SELECT * FROM test".foldWhile(List[(String,Int)]())({ (l, row) => 
+          (l :+ (row[String]("foo") -> row[Int]("bar"))) -> true 
+        }) aka "tuple stream" must_== Right(List("row1" -> 100, "row2" -> 200))
+      }
+
+    "handle failure" in withQueryResult(
+      rowList1(classOf[String] -> "foo") :+ "A" :+ "B") { implicit c =>
+      var i = 0
+        SQL"SELECT str".foldWhile(Set[String]()) { (l, row) => 
+      if (i == 0) { i = i+1; (l + row[String]("foo")) -> true } 
+      else sys.error("Failure")
+
+      } aka "aggregate on failure" must beLike {
+        case Left(err :: Nil) => err.getMessage aka "failure" must_== "Failure"
+      } and (i aka "row count" must_== 1)
+    }
+
+    "stop after first row" in withQueryResult(
+      rowList1(classOf[String] -> "foo") :+ "A" :+ "B") { implicit c =>
+      var i = 0
+        SQL"SELECT str".foldWhile(Set[String]()) { (l, row) => 
+      if (i == 0) { i = i+1; (l + row[String]("foo")) -> true } else (l, false)
+
+      } aka "partial aggregate" must_== Right(Set("A"))
+    }
+  }
+
   "Insertion" should {
     lazy implicit val con = connection(handleStatement withUpdateHandler {
       case UpdateExecution("INSERT ?", ExecutedParameter(1) :: Nil) => 1

framework/src/anorm/src/test/scala/anorm/SqlResultSpec.scala

@@ -203,6 +203,116 @@ object SqlResultSpec extends org.specs2.mutable.Specification with H2Database {
     }
   }
 
+  "Aggregation over all rows" should {
+    "release resources" in withQueryResult(
+      stringList :+ "A" :+ "B" :+ "C") { implicit c =>
+
+      val res: SqlQueryResult = SQL"SELECT str".executeQuery()
+      var closed = false
+      val probe = resource.managed(
+        new java.io.Closeable { def close() = closed = true })
+
+      var i = 0      
+      lazy val agg = res.copy(resultSet = 
+        res.resultSet.and(probe).map(_._1)).fold(List[Int]()) { 
+        (l, _) => i = i+1; l :+ i
+      }
+
+      agg aka "aggregation" must_== Right(List(1, 2, 3)) and (
+        closed aka "resource release" must beTrue) and (
+        i aka "row count" must_== 3)
+
+    }
+
+    "release resources on exception" in withQueryResult(
+      stringList :+ "A" :+ "B" :+ "C") { implicit c =>
+
+      val res: SqlQueryResult = SQL"SELECT str".executeQuery()
+      var closed = false
+      val probe = resource.managed(
+        new java.io.Closeable { def close() = closed = true })
+
+      var i = 0      
+      lazy val agg = res.copy(resultSet = res.resultSet.and(probe).map(_._1)).
+        fold(List[Int]()) { (l, _) => 
+          if (i == 1) sys.error("Unexpected") else { i = i +1; l :+ i }
+        } 
+
+      agg aka "aggregation" must beLike {
+        case Left(err :: Nil) =>
+          err.getMessage aka "failure" must_== "Unexpected"
+      } and (closed aka("resource release") must beTrue) and (
+        i aka "row count" must_== 1)
+
+    }
+  }
+
+  "Aggregation over variable number of rows" should {
+    "release resources" in withQueryResult(
+      stringList :+ "A" :+ "B" :+ "C") { implicit c =>
+
+      val res: SqlQueryResult = SQL"SELECT str".executeQuery()
+      var closed = false
+      val probe = resource.managed(
+        new java.io.Closeable { def close() = closed = true })
+
+      var i = 0      
+      lazy val agg = res.copy(resultSet = 
+        res.resultSet.and(probe).map(_._1)).foldWhile(List[Int]()) { 
+        (l, _) => i = i+1; (l :+ i) -> true
+      }
+
+      agg aka "aggregation" must_== Right(List(1, 2, 3)) and (
+        closed aka "resource release" must beTrue) and (
+        i aka "row count" must_== 3)
+
+    }
+
+    "release resources on exception" in withQueryResult(
+      stringList :+ "A" :+ "B" :+ "C") { implicit c =>
+
+      val res: SqlQueryResult = SQL"SELECT str".executeQuery()
+      var closed = false
+      val probe = resource.managed(
+        new java.io.Closeable { def close() = closed = true })
+
+      var i = 0      
+      lazy val agg = res.copy(resultSet = res.resultSet.and(probe).map(_._1)).
+        foldWhile(List[Int]()) { (l, _) => 
+          if (i == 1) sys.error("Unexpected") else { 
+            i = i +1; (l :+ i) -> true 
+          }
+        } 
+
+      agg aka "aggregation" must beLike {
+        case Left(err :: Nil) =>
+          err.getMessage aka "failure" must_== "Unexpected"
+      } and (closed aka "resource release" must beTrue) and (
+        i aka "row count" must_== 1)
+
+    }
+
+    "stop after second row & release resources" in withQueryResult(
+      stringList :+ "A" :+ "B" :+ "C") { implicit c =>
+
+      val res: SqlQueryResult = SQL"SELECT str".executeQuery()
+      var closed = false
+      val probe = resource.managed(
+        new java.io.Closeable { def close() = closed = true })
+
+      var i = 0      
+      lazy val agg = res.copy(resultSet = res.resultSet.and(probe).map(_._1)).
+        foldWhile(List[Int]()) { (l, _) => 
+          if (i == 2) (l, false) else { i = i +1; (l :+ i) -> true }
+        } 
+
+      agg aka "aggregation" must_== Right(List(1, 2)) and (
+        closed aka "resource release" must beTrue) and (
+        i aka "row count" must_== 2)
+
+    }
+  }
+
   "SQL warning" should {
     "not be there on success" in withQueryResult(stringList :+ "A") { 
       implicit c =>