Fix issue2669: bsddb simple/legacy interface iteration silently fails

when database changes size during iteration.

It now behaves like a dictionary, the next attempt to get a value from
the iterator after the database has changed size will raise a RuntimeError.

2 files changed, 128 insertions, 75 deletions

diff --git a/Lib/bsddb/__init__.py b/Lib/bsddb/__init__.py
index 44cfdfb..70b9ccf 100644
--- a/Lib/bsddb/__init__.py
+++ b/Lib/bsddb/__init__.py
@@ -33,7 +33,7 @@
 #----------------------------------------------------------------------
 
 
-"""Support for Berkeley DB 3.3 through 4.6 with a simple interface.
+"""Support for Berkeley DB 4.x with a simple interface.
 
 For the full featured object oriented interface use the bsddb.db module
 instead.  It mirrors the Oracle Berkeley DB C API.
@@ -66,13 +66,8 @@ error = db.DBError  # So bsddb.error will mean something...
 
 import sys, os
 
-# for backwards compatibility with python versions older than 2.3, the
-# iterator interface is dynamically defined and added using a mixin
-# class.  old python can't tokenize it due to the yield keyword.
-if sys.version >= '2.3':
-    import UserDict
-    from weakref import ref
-    exec """
+import UserDict
+from weakref import ref
 class _iter_mixin(UserDict.DictMixin):
     def _make_iter_cursor(self):
         cur = _DeadlockWrap(self.db.cursor)
@@ -87,67 +82,80 @@ class _iter_mixin(UserDict.DictMixin):
         return lambda ref: self._cursor_refs.pop(key, None)
 
     def __iter__(self):
+        self._kill_iteration = False
+        self._in_iter += 1
         try:
-            cur = self._make_iter_cursor()
-
-            # FIXME-20031102-greg: race condition.  cursor could
-            # be closed by another thread before this call.
-
-            # since we're only returning keys, we call the cursor
-            # methods with flags=0, dlen=0, dofs=0
-            key = _DeadlockWrap(cur.first, 0,0,0)[0]
-            yield key
-
-            next = cur.next
-            while 1:
-                try:
-                    key = _DeadlockWrap(next, 0,0,0)[0]
-                    yield key
-                except _bsddb.DBCursorClosedError:
-                    cur = self._make_iter_cursor()
-                    # FIXME-20031101-greg: race condition.  cursor could
-                    # be closed by another thread before this call.
-                    _DeadlockWrap(cur.set, key,0,0,0)
-                    next = cur.next
-        except _bsddb.DBNotFoundError:
-            return
-        except _bsddb.DBCursorClosedError:
-            # the database was modified during iteration.  abort.
-            return
+            try:
+                cur = self._make_iter_cursor()
+
+                # FIXME-20031102-greg: race condition.  cursor could
+                # be closed by another thread before this call.
+
+                # since we're only returning keys, we call the cursor
+                # methods with flags=0, dlen=0, dofs=0
+                key = _DeadlockWrap(cur.first, 0,0,0)[0]
+                yield key
+
+                next = cur.next
+                while 1:
+                    try:
+                        key = _DeadlockWrap(next, 0,0,0)[0]
+                        yield key
+                    except _bsddb.DBCursorClosedError:
+                        if self._kill_iteration:
+                            raise RuntimeError('Database changed size '
+                                               'during iteration.')
+                        cur = self._make_iter_cursor()
+                        # FIXME-20031101-greg: race condition.  cursor could
+                        # be closed by another thread before this call.
+                        _DeadlockWrap(cur.set, key,0,0,0)
+                        next = cur.next
+            except _bsddb.DBNotFoundError:
+                pass
+            except _bsddb.DBCursorClosedError:
+                # the database was modified during iteration.  abort.
+                pass
+        finally:
+            self._in_iter -= 1
 
     def iteritems(self):
         if not self.db:
             return
+        self._kill_iteration = False
+        self._in_iter += 1
         try:
-            cur = self._make_iter_cursor()
-
-            # FIXME-20031102-greg: race condition.  cursor could
-            # be closed by another thread before this call.
-
-            kv = _DeadlockWrap(cur.first)
-            key = kv[0]
-            yield kv
-
-            next = cur.next
-            while 1:
-                try:
-                    kv = _DeadlockWrap(next)
-                    key = kv[0]
-                    yield kv
-                except _bsddb.DBCursorClosedError:
-                    cur = self._make_iter_cursor()
-                    # FIXME-20031101-greg: race condition.  cursor could
-                    # be closed by another thread before this call.
-                    _DeadlockWrap(cur.set, key,0,0,0)
-                    next = cur.next
-        except _bsddb.DBNotFoundError:
-            return
-        except _bsddb.DBCursorClosedError:
-            # the database was modified during iteration.  abort.
-            return
-"""
-else:
-    class _iter_mixin: pass
+            try:
+                cur = self._make_iter_cursor()
+
+                # FIXME-20031102-greg: race condition.  cursor could
+                # be closed by another thread before this call.
+
+                kv = _DeadlockWrap(cur.first)
+                key = kv[0]
+                yield kv
+
+                next = cur.next
+                while 1:
+                    try:
+                        kv = _DeadlockWrap(next)
+                        key = kv[0]
+                        yield kv
+                    except _bsddb.DBCursorClosedError:
+                        if self._kill_iteration:
+                            raise RuntimeError('Database changed size '
+                                               'during iteration.')
+                        cur = self._make_iter_cursor()
+                        # FIXME-20031101-greg: race condition.  cursor could
+                        # be closed by another thread before this call.
+                        _DeadlockWrap(cur.set, key,0,0,0)
+                        next = cur.next
+            except _bsddb.DBNotFoundError:
+                pass
+            except _bsddb.DBCursorClosedError:
+                # the database was modified during iteration.  abort.
+                pass
+        finally:
+            self._in_iter -= 1
 
 
 class _DBWithCursor(_iter_mixin):
@@ -176,6 +184,8 @@ class _DBWithCursor(_iter_mixin):
         # a collection of all DBCursor objects currently allocated
         # by the _iter_mixin interface.
         self._cursor_refs = {}
+        self._in_iter = 0
+        self._kill_iteration = False
 
     def __del__(self):
         self.close()
@@ -225,6 +235,8 @@ class _DBWithCursor(_iter_mixin):
     def __setitem__(self, key, value):
         self._checkOpen()
         self._closeCursors()
+        if self._in_iter and key not in self:
+            self._kill_iteration = True
         def wrapF():
             self.db[key] = value
         _DeadlockWrap(wrapF)  # self.db[key] = value
@@ -232,6 +244,8 @@ class _DBWithCursor(_iter_mixin):
     def __delitem__(self, key):
         self._checkOpen()
         self._closeCursors()
+        if self._in_iter and key in self:
+            self._kill_iteration = True
         def wrapF():
             del self.db[key]
         _DeadlockWrap(wrapF)  # del self.db[key]
diff --git a/Lib/test/test_bsddb.py b/Lib/test/test_bsddb.py
index acd4972..f6e6ad0 100755
--- a/Lib/test/test_bsddb.py
+++ b/Lib/test/test_bsddb.py
@@ -66,9 +66,6 @@ class TestBSDDB(unittest.TestCase):
         self.assertSetEquals(d.iteritems(), f.iteritems())
 
     def test_iter_while_modifying_values(self):
-        if not hasattr(self.f, '__iter__'):
-            return
-
         di = iter(self.d)
         while 1:
             try:
@@ -80,20 +77,62 @@ class TestBSDDB(unittest.TestCase):
         # it should behave the same as a dict.  modifying values
         # of existing keys should not break iteration.  (adding
         # or removing keys should)
+        loops_left = len(self.f)
         fi = iter(self.f)
         while 1:
             try:
                 key = fi.next()
                 self.f[key] = 'modified '+key
+                loops_left -= 1
             except StopIteration:
                 break
+        self.assertEqual(loops_left, 0)
 
         self.test_mapping_iteration_methods()
 
-    def test_iteritems_while_modifying_values(self):
-        if not hasattr(self.f, 'iteritems'):
-            return
+    def test_iter_abort_on_changed_size(self):
+        def DictIterAbort():
+            di = iter(self.d)
+            while 1:
+                try:
+                    di.next()
+                    self.d['newkey'] = 'SPAM'
+                except StopIteration:
+                    break
+        self.assertRaises(RuntimeError, DictIterAbort)
+
+        def DbIterAbort():
+            fi = iter(self.f)
+            while 1:
+                try:
+                    fi.next()
+                    self.f['newkey'] = 'SPAM'
+                except StopIteration:
+                    break
+        self.assertRaises(RuntimeError, DbIterAbort)
 
+    def test_iteritems_abort_on_changed_size(self):
+        def DictIteritemsAbort():
+            di = self.d.iteritems()
+            while 1:
+                try:
+                    di.next()
+                    self.d['newkey'] = 'SPAM'
+                except StopIteration:
+                    break
+        self.assertRaises(RuntimeError, DictIteritemsAbort)
+
+        def DbIteritemsAbort():
+            fi = self.f.iteritems()
+            while 1:
+                try:
+                    key, value = fi.next()
+                    del self.f[key]
+                except StopIteration:
+                    break
+        self.assertRaises(RuntimeError, DbIteritemsAbort)
+
+    def test_iteritems_while_modifying_values(self):
         di = self.d.iteritems()
         while 1:
             try:
@@ -105,13 +144,16 @@ class TestBSDDB(unittest.TestCase):
         # it should behave the same as a dict.  modifying values
         # of existing keys should not break iteration.  (adding
         # or removing keys should)
+        loops_left = len(self.f)
         fi = self.f.iteritems()
         while 1:
             try:
                 k, v = fi.next()
                 self.f[k] = 'modified '+v
+                loops_left -= 1
             except StopIteration:
                 break
+        self.assertEqual(loops_left, 0)
 
         self.test_mapping_iteration_methods()
 
@@ -177,8 +219,8 @@ class TestBSDDB(unittest.TestCase):
         # the database write and locking+threading support is enabled
         # the cursor's read lock will deadlock the write lock request..
 
-        # test the iterator interface (if present)
-        if hasattr(self.f, 'iteritems'):
+        # test the iterator interface
+        if True:
             if debug: print "D"
             i = self.f.iteritems()
             k,v = i.next()
@@ -213,10 +255,7 @@ class TestBSDDB(unittest.TestCase):
         self.assert_(self.f[k], "be gone with ye deadlocks")
 
     def test_for_cursor_memleak(self):
-        if not hasattr(self.f, 'iteritems'):
-            return
-
-        # do the bsddb._DBWithCursor _iter_mixin internals leak cursors?
+        # do the bsddb._DBWithCursor iterator internals leak cursors?
         nc1 = len(self.f._cursor_refs)
         # create iterator
         i = self.f.iteritems()