Ensure atomic_batch works over DBProxy

Fixes #1494

Author: cburgdorf

tests/trinity/core/database/test_database_over_ipc_manager.py

@@ -94,3 +94,19 @@ def test_db_over_ipc_manager(manager):
 
     with pytest.raises(KeyError):
         db[b'not-present']
+
+
+def test_atomic_batch_over_ipc_manager(manager):
+    db = manager.get_db()
+
+    assert b'key-b' not in db
+    assert b'key-c' not in db
+
+    with db.atomic_batch() as batch:
+        batch.set(b'key-b', b'value-b')
+        batch.set(b'key-c', b'value-c')
+
+    assert b'key-b' in db
+    assert db[b'key-b'] == b'value-b'
+    assert b'key-c' in db
+    assert db[b'key-c'] == b'value-c'

trinity/db/base.py

@@ -1,10 +1,15 @@
+from contextlib import contextmanager
+from typing import (
+    Generator,
+)
 # Typeshed definitions for multiprocessing.managers is incomplete, so ignore them for now:
 # https://github.com/python/typeshed/blob/85a788dbcaa5e9e9a62e55f15d44530cd28ba830/stdlib/3/multiprocessing/managers.pyi#L3
 from multiprocessing.managers import (  # type: ignore
     BaseProxy,
 )
 
 from eth.db.backends.base import BaseDB
+from eth.db.atomic import AtomicDBWriteBatch
 
 from trinity.utils.mp import async_method
 
@@ -15,12 +20,13 @@ class DBProxy(BaseProxy):
         '__delitem__',
         '__getitem__',
         '__setitem__',
+        'atomic_batch',
+        'coro_set',
+        'coro_exists',
         'delete',
         'exists',
         'get',
         'set',
-        'coro_set',
-        'coro_exists',
     )
     coro_set = async_method('set')
     coro_exists = async_method('exists')
@@ -49,6 +55,11 @@ def exists(self, key: bytes) -> bool:
     def __contains__(self, key: bytes) -> bool:
         return self._callmethod('__contains__', (key,))
 
+    @contextmanager
+    def atomic_batch(self) -> Generator['AtomicDBWriteBatch', None, None]:
+        with AtomicDBWriteBatch._commit_unless_raises(self) as readable_batch:
+            yield readable_batch
+
 
 class AsyncBaseDB(BaseDB):