Iterator streaming for Query and Returning

magnum/src/main/scala/com/augustnagro/magnum/Query.scala

@@ -1,5 +1,6 @@
 package com.augustnagro.magnum
 
+import scala.util.Using.Manager
 import scala.util.{Failure, Success, Using}
 
 case class Query[E](frag: Frag, reader: DbCodec[E]):
@@ -14,3 +15,20 @@ case class Query[E](frag: Frag, reader: DbCodec[E]):
     ) match
       case Success(res) => res
       case Failure(t)   => throw SqlException(frag, t)
+
+  /** Streaming [[Iterator]]. Set [[fetchSize]] to give the JDBC driver a hint
+    * as to how many rows to fetch per request
+    */
+  def iterator(
+      fetchSize: Int = 0
+  )(using con: DbCon, use: Manager): Iterator[E] =
+    logSql(frag)
+    try
+      val ps = use(con.connection.prepareStatement(frag.sqlString))
+      ps.setFetchSize(fetchSize)
+      frag.writer.write(ps, 1)
+      val rs = use(ps.executeQuery())
+      ResultSetIterator(rs, frag, reader)
+    catch case t => throw SqlException(frag, t)
+
+end Query

magnum/src/main/scala/com/augustnagro/magnum/ResultSetIterator.scala

@@ -0,0 +1,25 @@
+package com.augustnagro.magnum
+
+import java.sql.ResultSet
+
+private class ResultSetIterator[E](
+    rs: ResultSet,
+    frag: Frag,
+    reader: DbCodec[E]
+) extends Iterator[E] {
+
+  private var rsHasNext: Boolean =
+    try rs.next()
+    catch case t => throw SqlException(frag, t)
+
+  override def hasNext: Boolean = rsHasNext
+
+  override def next(): E =
+    if !rsHasNext then throw IllegalStateException("ResultSet is empty")
+    try
+      val e = reader.readSingle(rs)
+      rsHasNext = rs.next()
+      e
+    catch case t => throw SqlException(frag, t)
+
+}

magnum/src/main/scala/com/augustnagro/magnum/Returning.scala

@@ -21,3 +21,24 @@ case class Returning[E](frag: Frag, reader: DbCodec[E]):
     ) match
       case Success(res) => res
       case Failure(t)   => throw SqlException(frag, t)
+
+  /** Streaming [[Iterator]]. Set [[fetchSize]] to give the JDBC driver a hint
+    * as to how many rows to fetch per request
+    */
+  def iterator(
+      fetchSize: Int = 0
+  )(using con: DbCon, use: Manager): Iterator[E] =
+    logSql(frag)
+    try
+      val ps = use(con.connection.prepareStatement(frag.sqlString))
+      ps.setFetchSize(fetchSize)
+      frag.writer.write(ps, 1)
+      val hasResults = ps.execute()
+      if hasResults then
+        val rs = use(ps.getResultSet())
+        ResultSetIterator(rs, frag, reader)
+      else
+        throw UnsupportedOperationException("No results for RETURNING clause")
+    catch case t => throw SqlException(frag, t)
+
+end Returning

magnum/src/test/scala/PgTests.scala

@@ -420,6 +420,22 @@ class PgTests extends FunSuite, TestContainersFixtures:
       customPersonRepo.count
     assertEquals(count, 8L)
 
+  test(".query iterator"):
+    connect(ds()):
+      Using.Manager(implicit use =>
+        val it = sql"SELECT * FROM person".query[Person].iterator()
+        assertEquals(it.map(_.id).size, 8)
+      )
+
+  test(".returning iterator"):
+    connect(ds()):
+      Using.Manager(implicit use =>
+        val it = sql"UPDATE person set is_admin = false RETURNING first_name"
+          .returning[String]
+          .iterator()
+        assertEquals(it.size, 8)
+      )
+
   @Table(PostgresDbType, SqlNameMapper.CamelToSnakeCase)
   case class BigDec(id: Int, myBigDec: Option[BigDecimal]) derives DbCodec
 

project/build.properties

@@ -1 +1 @@
-sbt.version=1.8.2
+sbt.version=1.9.9