[PATCH 09/14] pytdb: Add support for tdb_traverse & tdb_traverse_read

[Kirill Smelkov]

From: Kirill Smelkov <kirr at mns.spb.ru>

This functions are sometimes needed for making more controlled db
traversal compared to firstkey/nextkey, and also they are used in
tdbtorture, so add support for them.

Signed-off-by: Kirill Smelkov <kirr at mns.spb.ru>
---
 lib/tdb/pytdb.c                |   81 ++++++++++++++++++++++++++++++++++++++++
 lib/tdb/python/tests/simple.py |   38 +++++++++++++++++++
 2 files changed, 119 insertions(+), 0 deletions(-)

diff --git a/lib/tdb/pytdb.c b/lib/tdb/pytdb.c
index 2f54fcf..6d74b39 100644
--- a/lib/tdb/pytdb.c
+++ b/lib/tdb/pytdb.c
@@ -456,6 +456,81 @@ static PyObject *obj_increment_seqnum_nonblock(PyTdbObject *self)
 	Py_RETURN_NONE;
 }
 
+/* traverse */
+struct pytraverse_state {
+	PyTdbObject *self;
+	PyObject *pyfn;
+};
+
+static int pytraverse_helper(TDB_CONTEXT *tdb, TDB_DATA key, TDB_DATA dbuf, void *state)
+{
+	struct pytraverse_state *pystate = state;
+
+	PyObject *pykey=NULL, *pydata=NULL;
+	PyObject *pyret=NULL;
+	int ret=-1;
+
+	pykey = __PyString_FromTDB_DATA(key, /*release=*/false);
+	if (!pykey)
+		goto out;
+
+	pydata = __PyString_FromTDB_DATA(dbuf, /*release=*/false);
+	if (!pydata)
+		goto out;
+
+	pyret = PyObject_CallFunctionObjArgs(pystate->pyfn, pystate->self, pykey, pydata, NULL);
+	if (!pyret)
+		goto out;
+
+	if (PyObject_IsTrue(pyret))
+		ret = 0;	/* tell tdb to continue traversal */
+
+out:
+	Py_XDECREF(pykey);
+	Py_XDECREF(pydata);
+	Py_XDECREF(pyret);
+
+	return ret;
+}
+
+static PyObject *__obj_traverse(PyTdbObject *self, PyObject *args,
+	int (*traverse)(struct tdb_context *, tdb_traverse_func, void *))
+{
+	struct pytraverse_state st;
+	int ret;
+
+	st.self = self;
+	if (!PyArg_ParseTuple(args, "O", &st.pyfn))
+		return NULL;
+
+	ret = traverse(self->ctx,
+			(st.pyfn != Py_None ? pytraverse_helper : NULL),
+			&st);
+
+	/* check for py exceptions raised from-inside pytraverse_helper, and
+	 * re-raise it if needed */
+	if (PyErr_Occurred())
+		return NULL;
+
+	if (ret == -1) {
+		PyErr_SetTDBError(self->ctx);
+		return NULL;
+	}
+
+	return PyInt_FromLong(ret);
+}
+
+static PyObject *obj_traverse(PyTdbObject *self, PyObject *args)
+{
+	return __obj_traverse(self, args, tdb_traverse);
+}
+
+static PyObject *obj_traverse_read(PyTdbObject *self, PyObject *args)
+{
+	return __obj_traverse(self, args, tdb_traverse_read);
+}
+
+
 static PyMethodDef tdb_object_methods[] = {
 	{ "transaction_cancel", (PyCFunction)obj_transaction_cancel, METH_NOARGS, 
 		"S.transaction_cancel() -> None\n"
@@ -519,6 +594,12 @@ static PyMethodDef tdb_object_methods[] = {
 		"S.enable_seqnum() -> None" },
 	{ "increment_seqnum_nonblock", (PyCFunction)obj_increment_seqnum_nonblock, METH_NOARGS,
 		"S.increment_seqnum_nonblock() -> None" },
+	{ "traverse", (PyCFunction)obj_traverse, METH_VARARGS,
+		"S.traverse(fn|None) -> ntraversed\n"
+		"fn(key, data) -> continue?" },
+	{ "traverse_read", (PyCFunction)obj_traverse_read, METH_VARARGS,
+		"S.traverse_read(fn|None) -> ntraversed\n"
+		"fn(key, data) -> continue?" },
 	{ NULL }
 };
 
diff --git a/lib/tdb/python/tests/simple.py b/lib/tdb/python/tests/simple.py
index 4fe5008..cd65b22 100644
--- a/lib/tdb/python/tests/simple.py
+++ b/lib/tdb/python/tests/simple.py
@@ -170,6 +170,44 @@ class SimpleTdbTests(TestCase):
         self.tdb.transaction_commit()
         self.assertEquals("1", self.tdb["bloe"])
 
+    def test_traverse(self):
+        self.tdb["a"] = "1"
+        self.tdb["b"] = "2"
+        self.tdb["c"] = "3"
+
+        n = self.tdb.traverse(None)
+        self.assertEquals(n, 3)
+
+        v = []
+        def fn(db, key, data):
+            v.append( (key, data) )
+            return True
+        n = self.tdb.traverse(fn)
+        self.assertEquals(n, 3)
+        self.assertEquals(set(v), set([("a","1"), ("b","2"), ("c","3")]))
+
+        v = []
+        def fn(db, key, data):
+            return False
+        n = self.tdb.traverse(fn)
+        self.assertEquals(n, 1)
+
+        v = []
+        def fn(db, key, data):
+            1/0
+        self.assertRaises(ZeroDivisionError, self.tdb.traverse, fn)
+
+    def test_traverse_read(self):
+        self.tdb["a"] = "1"
+        self.tdb["b"] = "2"
+        self.tdb["c"] = "3"
+
+        n = self.tdb.traverse_read(None)
+        self.assertEquals(n, 3)
+
+        # the rest is the same as in .traverse
+
+
     def test_iterkeys(self):
         self.tdb["bloe"] = "2"
         self.tdb["bla"] = "25"
-- 
1.7.3.1.50.g1e633