Include missing last set in cursor - fixes porsager#150 porsager#81

porsager · porsager · commit cea5f73121ca · 2021-06-18T08:53:08.000+02:00
diff --git a/lib/connection.js b/lib/connection.js
@@ -61,11 +61,11 @@ function Connection(options = {}) {
     error
   })
 
-  function onsuspended(x) {
+  function onsuspended(x, done) {
     new Promise(r => r(backend.query.cursor(
       backend.query.cursor.rows === 1 ? x[0] : x
     ))).then(x => {
-      x === END
+      x === END || done
         ? socket.write(frontend.Close())
         : socket.write(frontend.Execute(backend.query.cursor.rows))
     }).catch(err => {
@@ -75,7 +75,7 @@ function Connection(options = {}) {
   }
 
   function oncomplete() {
-    backend.query.cursor && socket.write(frontend.Close())
+    backend.query.cursor && onsuspended(backend.query.result, true)
   }
 
   function onparse() {
diff --git a/tests/index.js b/tests/index.js
@@ -956,6 +956,22 @@ t('Cursor custom n works', async() => {
   return ['10,10', order.join(',')]
 })
 
+ot('Cursor custom with rest n works', async() => {
+  const order = []
+  await sql`select * from generate_series(1,20)`.cursor(11, async(x) => {
+    order.push(x.length)
+  })
+  return ['11,9', order.join(',')]
+})
+
+ot('Cursor custom with less results than batch size works', async() => {
+  const order = []
+  await sql`select * from generate_series(1,20)`.cursor(21, async(x) => {
+    order.push(x.length)
+  })
+  return ['11,9', order.join(',')]
+})
+
 t('Cursor cancel works', async() => {
   let result
   await sql`select * from generate_series(1,10) as x`.cursor(async({ x }) => {
