Support for await...of db.iterator() (#379)

README.md

@@ -311,32 +311,54 @@ A reference to the `db` that created this chained batch.
 
 An iterator allows you to _iterate_ the entire store or a range. It operates on a snapshot of the store, created at the time `db.iterator()` was called. This means reads on the iterator are unaffected by simultaneous writes. Most but not all implementations can offer this guarantee.
 
-An iterator keeps track of when a `next()` is in progress and when an `end()` has been called so it doesn't allow concurrent `next()` calls, it does allow `end()` while a `next()` is in progress and it doesn't allow either `next()` or `end()` after `end()` has been called.
-
-#### `iterator.next(callback)`
-
-Advance the iterator and yield the entry at that key. If an error occurs, the `callback` function will be called with an `Error`. Otherwise, the `callback` receives `null`, a `key` and a `value`. The type of `key` and `value` depends on the options passed to `db.iterator()`.
+Iterators can be consumed with [`for await...of`](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Statements/for-await...of) or by manually calling `iterator.next()` in succession. In the latter mode, `iterator.end()` must always be called. In contrast, finishing, throwing or breaking from a `for await...of` loop automatically calls `iterator.end()`.
 
-If the iterator has reached its end, both `key` and `value` will be `undefined`. This happens in the following situations:
+An iterator reaches its natural end in the following situations:
 
 - The end of the store has been reached
 - The end of the range has been reached
 - The last `iterator.seek()` was out of range.
 
-**Note:** Don't forget to call `iterator.end()`, even if you received an error.
+An iterator keeps track of when a `next()` is in progress and when an `end()` has been called so it doesn't allow concurrent `next()` calls, it does allow `end()` while a `next()` is in progress and it doesn't allow either `next()` or `end()` after `end()` has been called.
+
+#### `for await...of iterator`
+
+Yields arrays containing a `key` and `value`. The type of `key` and `value` depends on the options passed to `db.iterator()`.
+
+```js
+try {
+  for await (const [key, value] of db.iterator()) {
+    console.log(key)
+  }
+} catch (err) {
+  console.error(err)
+}
+```
+
+Note for implementors: this uses `iterator.next()` and `iterator.end()` under the hood so no further method implementations are needed to support `for await...of`.
+
+#### `iterator.next([callback])`
+
+Advance the iterator and yield the entry at that key. If an error occurs, the `callback` function will be called with an `Error`. Otherwise, the `callback` receives `null`, a `key` and a `value`. The type of `key` and `value` depends on the options passed to `db.iterator()`. If the iterator has reached its natural end, both `key` and `value` will be `undefined`.
+
+If no callback is provided, a promise is returned for either an array (containing a `key` and `value`) or `undefined` if the iterator reached its natural end.
+
+**Note:** Always call `iterator.end()`, even if you received an error and even if the iterator reached its natural end.
 
 #### `iterator.seek(target)`
 
-Seek the iterator to a given key or the closest key. Subsequent calls to `iterator.next()` will yield entries with keys equal to or larger than `target`, or equal to or smaller than `target` if the `reverse` option passed to `db.iterator()` was true.
+Seek the iterator to a given key or the closest key. Subsequent calls to `iterator.next()` (including implicit calls in a `for await...of` loop) will yield entries with keys equal to or larger than `target`, or equal to or smaller than `target` if the `reverse` option passed to `db.iterator()` was true.
 
-If range options like `gt` were passed to `db.iterator()` and `target` does not fall within that range, the iterator will reach its end.
+If range options like `gt` were passed to `db.iterator()` and `target` does not fall within that range, the iterator will reach its natural end.
 
 **Note:** At the time of writing, [`leveldown`][leveldown] is the only known implementation to support `seek()`. In other implementations, it is a noop.
 
-#### `iterator.end(callback)`
+#### `iterator.end([callback])`
 
 End iteration and free up underlying resources. The `callback` function will be called with no arguments on success or with an `Error` if ending failed for any reason.
 
+If no callback is provided, a promise is returned.
+
 #### `iterator.db`
 
 A reference to the `db` that created this iterator.

abstract-iterator.js

@@ -11,18 +11,29 @@ function AbstractIterator (db) {
 }
 
 AbstractIterator.prototype.next = function (callback) {
-  if (typeof callback !== 'function') {
+  // In callback mode, we return `this`
+  let ret = this
+
+  if (callback === undefined) {
+    ret = new Promise(function (resolve, reject) {
+      callback = function (err, key, value) {
+        if (err) reject(err)
+        else if (key === undefined && value === undefined) resolve()
+        else resolve([key, value])
+      }
+    })
+  } else if (typeof callback !== 'function') {
     throw new Error('next() requires a callback argument')
   }
 
   if (this._ended) {
     this._nextTick(callback, new Error('cannot call next() after end()'))
-    return this
+    return ret
   }
 
   if (this._nexting) {
     this._nextTick(callback, new Error('cannot call next() before previous next() has completed'))
-    return this
+    return ret
   }
 
   this._nexting = true
@@ -31,7 +42,7 @@ AbstractIterator.prototype.next = function (callback) {
     callback(err, ...rest)
   })
 
-  return this
+  return ret
 }
 
 AbstractIterator.prototype._next = function (callback) {
@@ -53,22 +64,46 @@ AbstractIterator.prototype.seek = function (target) {
 AbstractIterator.prototype._seek = function (target) {}
 
 AbstractIterator.prototype.end = function (callback) {
-  if (typeof callback !== 'function') {
+  let promise
+
+  if (callback === undefined) {
+    promise = new Promise(function (resolve, reject) {
+      callback = function (err) {
+        if (err) reject(err)
+        else resolve()
+      }
+    })
+  } else if (typeof callback !== 'function') {
     throw new Error('end() requires a callback argument')
   }
 
   if (this._ended) {
-    return this._nextTick(callback, new Error('end() already called on iterator'))
+    this._nextTick(callback, new Error('end() already called on iterator'))
+    return promise
   }
 
   this._ended = true
   this._end(callback)
+
+  return promise
 }
 
 AbstractIterator.prototype._end = function (callback) {
   this._nextTick(callback)
 }
 
+AbstractIterator.prototype[Symbol.asyncIterator] = async function * () {
+  try {
+    let kv
+
+    while ((kv = (await this.next())) !== undefined) {
+      yield kv
+    }
+  } finally {
+    if (!this._ended) await this.end()
+  }
+}
+
 // Expose browser-compatible nextTick for dependents
 AbstractIterator.prototype._nextTick = require('./next-tick')
 

test/async-iterator-test.js

@@ -0,0 +1,203 @@
+'use strict'
+
+const input = [{ key: '1', value: '1' }, { key: '2', value: '2' }]
+
+let db
+
+exports.setup = function (test, testCommon) {
+  test('setup', function (t) {
+    t.plan(2)
+
+    db = testCommon.factory()
+    db.open(function (err) {
+      t.ifError(err, 'no open() error')
+
+      db.batch(input.map(entry => ({ ...entry, type: 'put' })), function (err) {
+        t.ifError(err, 'no batch() error')
+      })
+    })
+  })
+}
+
+exports.asyncIterator = function (test, testCommon) {
+  test('for await...of db.iterator()', async function (t) {
+    t.plan(2)
+
+    const it = db.iterator({ keyAsBuffer: false, valueAsBuffer: false })
+    const output = []
+
+    for await (const [key, value] of it) {
+      output.push({ key, value })
+    }
+
+    t.ok(it._ended, 'ended')
+    t.same(output, input)
+  })
+
+  test('for await...of db.iterator() does not permit reuse', async function (t) {
+    t.plan(3)
+
+    const it = db.iterator()
+
+    // eslint-disable-next-line no-unused-vars
+    for await (const [key, value] of it) {
+      t.pass('nexted')
+    }
+
+    try {
+      // eslint-disable-next-line no-unused-vars
+      for await (const [key, value] of it) {
+        t.fail('should not be called')
+      }
+    } catch (err) {
+      t.is(err.message, 'cannot call next() after end()')
+    }
+  })
+
+  test('for await...of db.iterator() ends on user error', async function (t) {
+    t.plan(2)
+
+    const it = db.iterator()
+
+    try {
+      // eslint-disable-next-line no-unused-vars, no-unreachable-loop
+      for await (const kv of it) {
+        throw new Error('user error')
+      }
+    } catch (err) {
+      t.is(err.message, 'user error')
+      t.ok(it._ended, 'ended')
+    }
+  })
+
+  test('for await...of db.iterator() with user error and end() error', async function (t) {
+    t.plan(3)
+
+    const it = db.iterator()
+    const end = it._end
+
+    it._end = function (callback) {
+      end.call(this, function (err) {
+        t.ifError(err, 'no real error from end()')
+        callback(new Error('end error'))
+      })
+    }
+
+    try {
+      // eslint-disable-next-line no-unused-vars, no-unreachable-loop
+      for await (const kv of it) {
+        throw new Error('user error')
+      }
+    } catch (err) {
+      // TODO: ideally, this would be a combined aka aggregate error
+      t.is(err.message, 'user error')
+      t.ok(it._ended, 'ended')
+    }
+  })
+
+  test('for await...of db.iterator() ends on iterator error', async function (t) {
+    t.plan(3)
+
+    const it = db.iterator()
+
+    it._next = function (callback) {
+      t.pass('nexted')
+      this._nextTick(callback, new Error('iterator error'))
+    }
+
+    try {
+      // eslint-disable-next-line no-unused-vars
+      for await (const kv of it) {
+        t.fail('should not yield results')
+      }
+    } catch (err) {
+      t.is(err.message, 'iterator error')
+      t.ok(it._ended, 'ended')
+    }
+  })
+
+  test('for await...of db.iterator() with iterator error and end() error', async function (t) {
+    t.plan(4)
+
+    const it = db.iterator()
+    const end = it._end
+
+    it._next = function (callback) {
+      t.pass('nexted')
+      this._nextTick(callback, new Error('iterator error'))
+    }
+
+    it._end = function (callback) {
+      end.call(this, function (err) {
+        t.ifError(err, 'no real error from end()')
+        callback(new Error('end error'))
+      })
+    }
+
+    try {
+      // eslint-disable-next-line no-unused-vars
+      for await (const kv of it) {
+        t.fail('should not yield results')
+      }
+    } catch (err) {
+      // TODO: ideally, this would be a combined aka aggregate error
+      t.is(err.message, 'end error')
+      t.ok(it._ended, 'ended')
+    }
+  })
+
+  test('for await...of db.iterator() ends on user break', async function (t) {
+    t.plan(2)
+
+    const it = db.iterator()
+
+    // eslint-disable-next-line no-unused-vars, no-unreachable-loop
+    for await (const kv of it) {
+      t.pass('got a chance to break')
+      break
+    }
+
+    t.ok(it._ended, 'ended')
+  })
+
+  test('for await...of db.iterator() with user break and end() error', async function (t) {
+    t.plan(4)
+
+    const it = db.iterator()
+    const end = it._end
+
+    it._end = function (callback) {
+      end.call(this, function (err) {
+        t.ifError(err, 'no real error from end()')
+        callback(new Error('end error'))
+      })
+    }
+
+    try {
+      // eslint-disable-next-line no-unused-vars, no-unreachable-loop
+      for await (const kv of it) {
+        t.pass('got a chance to break')
+        break
+      }
+    } catch (err) {
+      t.is(err.message, 'end error')
+      t.ok(it._ended, 'ended')
+    }
+  })
+}
+
+exports.teardown = function (test, testCommon) {
+  test('teardown', function (t) {
+    t.plan(1)
+
+    db.close(function (err) {
+      t.ifError(err, 'no close() error')
+    })
+  })
+}
+
+exports.all = function (test, testCommon) {
+  exports.setup(test, testCommon)
+  exports.asyncIterator(test, testCommon)
+  exports.teardown(test, testCommon)
+}

test/index.js

@@ -31,6 +31,7 @@ function suite (options) {
 
   require('./iterator-test').all(test, testCommon)
   require('./iterator-range-test').all(test, testCommon)
+  require('./async-iterator-test').all(test, testCommon)
 
   if (testCommon.seek) {
     require('./iterator-seek-test').all(test, testCommon)