datastore.py: update to meet new datastore builtin PR.

The new datastore PR now has a generation count; this passes
the tests against that now.

Also copies tests from c-lightning, and adds ugprade test.

Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>

datastore/datastore-plugin.py

@@ -2,72 +2,104 @@
 """This does the actual datastore work, if the main plugin says there's no
 datastore support.  We can't even load this if there's real datastore support.
 """
-from pyln.client import Plugin, RpcError
+from pyln.client import Plugin, RpcException
+from collections import namedtuple
+import os
 import shelve
 
+# Error codes
+DATASTORE_DEL_DOES_NOT_EXIST = 1200
+DATASTORE_DEL_WRONG_GENERATION = 1201
+DATASTORE_UPDATE_ALREADY_EXISTS = 1202
+DATASTORE_UPDATE_DOES_NOT_EXIST = 1203
+DATASTORE_UPDATE_WRONG_GENERATION = 1204
 
 plugin = Plugin()
+Entry = namedtuple('Entry', ['generation', 'data'])
 
-
+def datastore_entry(key, entry: Entry):
-def datastore_entry(key, data):
     """Return a dict representing the entry"""
 
-    ret = {'key': key, 'hex': data.hex()}
+    # Entry may be a simple tuple; convert
+    entry = Entry(*entry)
+    ret = {'key': key, 'generation': entry.generation, 'hex': entry.data.hex()}
 
     # FFS, Python3 seems happy with \0 in UTF-8.
-    if 0 not in data:
+    if 0 not in entry.data:
         try:
-            ret['string'] = data.decode('utf8')
+            ret['string'] = entry.data.decode('utf8')
         except UnicodeDecodeError:
             pass
     return ret
 
 
 @plugin.method("datastore")
-def datastore(plugin, key, string=None, hex=None, mode="must-create"):
+def datastore(plugin, key, string=None, hex=None, mode="must-create", generation=None):
-    """Add a {key} and {hex}/{string} data to the data store"""
+    """Add/modify a {key} and {hex}/{string} data to the data store,
+optionally insisting it be {generation}"""
 
     if string is not None:
         if hex is not None:
-            raise RpcError("datastore", {'key': key},
+            raise RpcException("Cannot specify both string or hex")
-                           {'message': "Cannot specify both string or hex"})
         data = bytes(string, encoding="utf8")
     elif hex is None:
-        raise RpcError("datastore", {'key': key},
+        raise RpcException("Must specify string or hex")
-                       {'message': "Must specify string or hex"})
     else:
         data = bytes.fromhex(hex)
 
-    print("key={}, data={}, mode={}".format(key, data, mode))
     if mode == "must-create":
         if key in plugin.datastore:
-            raise RpcError("datastore", {'key': key},
+            raise RpcException("already exists", DATASTORE_UPDATE_ALREADY_EXISTS)
-                           {'message': "already exists"})
     elif mode == "must-replace":
         if key not in plugin.datastore:
-            raise RpcError("datastore", {'key': key},
+            raise RpcException("does not exist", DATASTORE_UPDATE_DOES_NOT_EXIST)
-                           {'message': "does not exist"})
     elif mode == "create-or-replace":
+        if generation is not None:
+            raise RpcException("generation only valid with"
+                               " must-create/must-replace")
         pass
     elif mode == "must-append":
+        if generation is not None:
+            raise RpcException("generation only valid with"
+                               " must-create/must-replace")
         if key not in plugin.datastore:
-            raise RpcError("datastore", {'key': key},
+            raise RpcException("does not exist", DATASTORE_UPDATE_DOES_NOT_EXIST)
-                           {'message': "does not exist"})
+        data = plugin.datastore[key].data + data
-        data = plugin.datastore[key] + data
     elif mode == "create-or-append":
-        data = plugin.datastore.get(key, bytes()) + data
+        if generation is not None:
+            raise RpcException("generation only valid with"
+                               " must-create/must-replace")
+        data = plugin.datastore.get(key, Entry(0, bytes())).data + data
     else:
-        raise RpcError("datastore", {'key': key}, {'message': "invalid mode"})
+        raise RpcException("invalid mode")
 
-    plugin.datastore[key] = data
+    if key in plugin.datastore:
-    return datastore_entry(key, data)
+        entry = plugin.datastore[key]
+        if generation is not None:
+            if entry.generation != generation:
+                raise RpcException("generation is different",
+                                   DATASTORE_UPDATE_WRONG_GENERATION)
+        gen = entry.generation + 1
+    else:
+        gen = 0
+
+    plugin.datastore[key] = Entry(gen, data)
+    return datastore_entry(key, plugin.datastore[key])
 
 
 @plugin.method("deldatastore")
-def deldatastore(plugin, key):
+def deldatastore(plugin, key, generation=None):
     """Remove a {key} from the data store"""
 
-    ret = datastore_entry(key, plugin.datastore[key])
+    if not key in plugin.datastore:
+        raise RpcException("does not exist", DATASTORE_DEL_DOES_NOT_EXIST)
+
+    entry = plugin.datastore[key]
+    if generation is not None and entry.generation != generation:
+        raise RpcException("generation is different",
+                           DATASTORE_DEL_WRONG_GENERATION)
+
+    ret = datastore_entry(key, entry)
     del plugin.datastore[key]
     return ret
 
@@ -76,16 +108,32 @@ def deldatastore(plugin, key):
 def listdatastore(plugin, key=None):
     """List datastore entries"""
     if key is None:
-        return {'datastore': [datastore_entry(k, d)
+        return {'datastore': [datastore_entry(k, e)
-                              for k, d in plugin.datastore.items()]}
+                              for k, e in plugin.datastore.items()]}
     if key in plugin.datastore:
         return {'datastore': [datastore_entry(key, plugin.datastore[key])]}
     return {'datastore': []}
 
 
+def upgrade_store(plugin):
+    """Initial version of this plugin had no generation numbers"""
+    try:
+        oldstore = shelve.open('datastore.dat', 'r')
+    except:
+        return
+    plugin.log("Upgrading store to have generation numbers", level='unusual')
+    datastore = shelve.open('datastore_v1.dat', 'c')
+    for k, d in oldstore.items():
+        datastore[k] = Entry(0, d)
+    oldstore.close()
+    datastore.close()
+    os.unlink('datastore.dat')
+
+
 @plugin.init()
 def init(options, configuration, plugin):
-    plugin.datastore = shelve.open('datastore.dat')
+    upgrade_store(plugin)
+    plugin.datastore = shelve.open('datastore_v1.dat')
 
 
 plugin.run()

datastore/datastore.py

@@ -2,6 +2,7 @@
 from pyln.client import Plugin, RpcError
 import shelve
 import os
+from collections import namedtuple
 
 
 plugin = Plugin()
@@ -10,12 +11,12 @@ plugin = Plugin()
 def unload_store(plugin):
     """When we have a real store, we transfer our contents into it"""
     try:
-        datastore = shelve.open('datastore.dat', 'r')
+        datastore = shelve.open('datastore_v1.dat', 'r')
     except:
         return
 
-    plugin.log("Emptying store into main store!", level='unusual')
+    plugin.log("Emptying store into main store (resetting generations!)", level='unusual')
-    for k, d in datastore.items():
+    for k, (g, d) in datastore.items():
         try:
             plugin.rpc.datastore(k, d.hex())
         except RpcError as e:
@@ -23,7 +24,7 @@ def unload_store(plugin):
                        level='broken')
     datastore.close()
     plugin.log("Erasing our store", level='unusual')
-    os.unlink('datastore.dat')
+    os.unlink('datastore_v1.dat')
 
 
 @plugin.init()

datastore/test_datastore.py

@@ -0,0 +1,154 @@
+import os
+import shelve
+import time
+from pyln.testing.fixtures import *  # noqa: F401,F403
+from pyln.client import RpcError
+from pyln.testing.utils import only_one, wait_for
+
+plugin_path = os.path.join(os.path.dirname(__file__), "datastore.py")
+
+# Test taken from lightning/tests/test_misc.py
+def test_datastore(node_factory):
+    l1 = node_factory.get_node(options={'plugin': plugin_path})
+    time.sleep(5)
+
+    # Starts empty
+    assert l1.rpc.listdatastore() == {'datastore': []}
+    assert l1.rpc.listdatastore('somekey') == {'datastore': []}
+
+    # Add entries.
+    somedata = b'somedata'.hex()
+    somedata_expect = {'key': 'somekey',
+                       'generation': 0,
+                       'hex': somedata,
+                       'string': 'somedata'}
+    assert l1.rpc.datastore(key='somekey', hex=somedata) == somedata_expect
+
+    assert l1.rpc.listdatastore() == {'datastore': [somedata_expect]}
+    assert l1.rpc.listdatastore('somekey') == {'datastore': [somedata_expect]}
+    assert l1.rpc.listdatastore('otherkey') == {'datastore': []}
+
+    # Cannot add by default.
+    with pytest.raises(RpcError, match='already exists'):
+        l1.rpc.datastore(key='somekey', hex=somedata)
+
+    with pytest.raises(RpcError, match='already exists'):
+        l1.rpc.datastore(key='somekey', hex=somedata, mode="must-create")
+
+    # But can insist on replace.
+    l1.rpc.datastore(key='somekey', hex=somedata[:-4], mode="must-replace")
+    assert only_one(l1.rpc.listdatastore('somekey')['datastore'])['hex'] == somedata[:-4]
+    # And append works.
+    l1.rpc.datastore(key='somekey', hex=somedata[-4:-2], mode="must-append")
+    assert only_one(l1.rpc.listdatastore('somekey')['datastore'])['hex'] == somedata[:-2]
+    l1.rpc.datastore(key='somekey', hex=somedata[-2:], mode="create-or-append")
+    assert only_one(l1.rpc.listdatastore('somekey')['datastore'])['hex'] == somedata
+
+    # Generation will have increased due to three ops above.
+    somedata_expect['generation'] += 3
+    assert l1.rpc.listdatastore() == {'datastore': [somedata_expect]}
+
+    # Can't replace or append non-existing records if we say not to
+    with pytest.raises(RpcError, match='does not exist'):
+        l1.rpc.datastore(key='otherkey', hex=somedata, mode="must-replace")
+
+    with pytest.raises(RpcError, match='does not exist'):
+        l1.rpc.datastore(key='otherkey', hex=somedata, mode="must-append")
+
+    otherdata = b'otherdata'.hex()
+    otherdata_expect = {'key': 'otherkey',
+                        'generation': 0,
+                        'hex': otherdata,
+                        'string': 'otherdata'}
+    assert l1.rpc.datastore(key='otherkey', string='otherdata', mode="create-or-append") == otherdata_expect
+
+    assert l1.rpc.listdatastore('somekey') == {'datastore': [somedata_expect]}
+    assert l1.rpc.listdatastore('otherkey') == {'datastore': [otherdata_expect]}
+    assert l1.rpc.listdatastore('badkey') == {'datastore': []}
+
+    ds = l1.rpc.listdatastore()
+    # Order is undefined!
+    assert (ds == {'datastore': [somedata_expect, otherdata_expect]}
+            or ds == {'datastore': [otherdata_expect, somedata_expect]})
+
+    assert l1.rpc.deldatastore('somekey') == somedata_expect
+    assert l1.rpc.listdatastore() == {'datastore': [otherdata_expect]}
+    assert l1.rpc.listdatastore('somekey') == {'datastore': []}
+    assert l1.rpc.listdatastore('otherkey') == {'datastore': [otherdata_expect]}
+    assert l1.rpc.listdatastore('badkey') == {'datastore': []}
+    assert l1.rpc.listdatastore() == {'datastore': [otherdata_expect]}
+
+    # if it's not a string, won't print
+    badstring_expect = {'key': 'badstring',
+                        'generation': 0,
+                        'hex': '00'}
+    assert l1.rpc.datastore(key='badstring', hex='00') == badstring_expect
+    assert l1.rpc.listdatastore('badstring') == {'datastore': [badstring_expect]}
+    assert l1.rpc.deldatastore('badstring') == badstring_expect
+
+    # It's persistent
+    l1.restart()
+
+    assert l1.rpc.listdatastore() == {'datastore': [otherdata_expect]}
+
+    # We can insist generation match on update.
+    with pytest.raises(RpcError, match='generation is different'):
+        l1.rpc.datastore(key='otherkey', hex='00', mode='must-replace',
+                         generation=otherdata_expect['generation'] + 1)
+
+    otherdata_expect['generation'] += 1
+    otherdata_expect['string'] += 'a'
+    otherdata_expect['hex'] += '61'
+    assert (l1.rpc.datastore(key='otherkey', string='otherdataa',
+                             mode='must-replace',
+                             generation=otherdata_expect['generation'] - 1)
+            == otherdata_expect)
+    assert l1.rpc.listdatastore() == {'datastore': [otherdata_expect]}
+
+    # We can insist generation match on delete.
+    with pytest.raises(RpcError, match='generation is different'):
+        l1.rpc.deldatastore(key='otherkey',
+                            generation=otherdata_expect['generation'] + 1)
+
+    assert (l1.rpc.deldatastore(key='otherkey',
+                                generation=otherdata_expect['generation'])
+            == otherdata_expect)
+    assert l1.rpc.listdatastore() == {'datastore': []}
+
+
+def test_upgrade(node_factory):
+    l1 = node_factory.get_node()
+
+    datastore = shelve.open(os.path.join(l1.daemon.lightning_dir, 'regtest', 'datastore.dat'), 'c')
+    datastore['foo'] = b'foodata'
+    datastore['bar'] = b'bardata'
+    datastore.close()
+
+    # This "fails" because it unloads itself.
+    try:
+        l1.rpc.plugin_start(plugin_path)
+    except RpcError:
+        pass
+
+    l1.daemon.wait_for_log('Upgrading store to have generation numbers')
+    wait_for(lambda: not os.path.exists(os.path.join(l1.daemon.lightning_dir,
+                                                     'regtest',
+                                                     'datastore.dat')))
+
+    vals = l1.rpc.listdatastore()['datastore']
+    assert (vals == [{'key': 'foo',
+                      'generation': 0,
+                      'hex': b'foodata'.hex(),
+                      'string': 'foodata'},
+                     {'key': 'bar',
+                      'generation': 0,
+                      'hex': b'bardata'.hex(),
+                      'string': 'bardata'}]
+            or vals == [{'key': 'bar',
+                         'generation': 0,
+                         'hex': b'bardata'.hex(),
+                         'string': 'bardata'},
+                        {'key': 'foo',
+                         'generation': 0,
+                         'hex': b'foodata'.hex(),
+                         'string': 'foodata'}])