From 8a1acc96728569b6d4616bf66e3f9dd2b7dd59a7 Mon Sep 17 00:00:00 2001 From: adam Date: Wed, 13 Feb 2013 18:07:15 +0700 Subject: [PATCH] #47 --- pyejdb/pyejdb/__init__.py | 113 +++++++++++++++++++++++++++++++++++++++------- pyejdb/setup.py | 3 +- pyejdb/src/pdbcursor.c | 82 ++++++++++++++++++++++++++++++++- pyejdb/src/pejdb.c | 64 ++++++++++++++++++++++++-- pyejdb/src/pyejdb.c | 20 ++++---- pyejdb/src/pyejdb.h | 4 ++ pyejdb/test/test_one.py | 25 +++++++++- 7 files changed, 275 insertions(+), 36 deletions(-) diff --git a/pyejdb/pyejdb/__init__.py b/pyejdb/pyejdb/__init__.py index 7465ddd..6585f25 100644 --- a/pyejdb/pyejdb/__init__.py +++ b/pyejdb/pyejdb/__init__.py @@ -1,9 +1,10 @@ import _pyejdb +from functools import lru_cache from pprint import pprint -from collections import OrderedDict as odict from pyejdb import bson from pyejdb.typecheck import * import re +import numbers __all__ = [ @@ -18,9 +19,21 @@ __all__ = [ "JBOTSYNC", "DEFAULT_OPEN_MODE", - "check_oid" + "JBQRYCOUNT", + + "check_oid", + "version", + "libejdb_version" ] +version_tuple = (1, 0, 1) + +def get_version_string(): + return '.'.join(map(str, version_tuple)) + +version = get_version_string() +libejdb_version = _pyejdb.ejdb_version() + #OID check RE _oidRE = re.compile("^[0-9a-f]{24}$") @@ -34,12 +47,66 @@ JBOLCKNB = _pyejdb.JBOLCKNB JBOTSYNC = _pyejdb.JBOTSYNC DEFAULT_OPEN_MODE = JBOWRITER | JBOCREAT | JBOTSYNC +#Query flags +JBQRYCOUNT = _pyejdb.JBQRYCOUNT +#Misc def check_oid(oid): if not isinstance(oid, str) or _oidRE.match(oid) is None: raise ValueError("Invalid OID: %s" % oid) +class EJDBCursorWrapper(object): + def __init__(self, dbcursor): + self.__pos = 0 + self.__cursor = dbcursor + self.__len = self.__cursor.length() + + def __del__(self): + self.__cursor.close() + + def __len__(self): + return self.__len + + def __getitem__(self, key): + if isinstance(key, slice): + return [self[idx] for idx in range(*key.indices(len(self)))] + elif isinstance(key, numbers.Number): + if key < 0: + key += len(self) + if key >= len(self): + raise IndexError("The index (%d) is out of range." % key) + return self.get(key) + else: + raise TypeError("Invalid argument type.") + + def __iter__(self): + return self + + def __next__(self): + if self.__pos >= self.__len: + raise StopIteration + self.__pos += 1 + return self.get(self.__pos - 1) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + def reset(self): + self.__pos = 0 + + @lru_cache(maxsize=8192) + def get(self, idx): + bsdata = self.__cursor.get(idx) + return bson.parse_bytes(bsdata) if bsdata is not None else None + + def close(self): + self.__cursor.close() + + class EJDB(object): def __init__(self, fpath, mode=DEFAULT_OPEN_MODE): #pprint (vars(_pyejdb)) @@ -80,23 +147,35 @@ class EJDB(object): def find(self, cname : str, qobj : optional(dict)=None, *args, **kwargs): if not qobj: qobj = {} qobj = bson.serialize_to_bytes(qobj) - hints = bson.serialize_to_bytes(kwargs["hints"] if "hints" in kwargs else {}) + hints = bson.serialize_to_bytes(kwargs.get("hints", {})) orarr = [bson.serialize_to_bytes(x) for x in args] - qflags = kwargs["qflags"] if "qflags" in kwargs else 0 - return self.__ejdb.find(cname, qobj, orarr, hints, qflags) - - - - - - - - - - - - + qflags = kwargs.get("qflags", 0) + cursor = self.__ejdb.find(cname, qobj, orarr, hints, qflags) + return cursor if isinstance(cursor, numbers.Number) else EJDBCursorWrapper(cursor) + + def findOne(self, cname : str, qobj : optional(dict)=None, *args, **kwargs): + hints = kwargs.get("hints", {}) + hints["$max"] = 1 + kwargs["hints"] = hints + with self.find(cname, qobj, *args, **kwargs) as res: + return res[0] if len(res) > 0 else None + + def update(self, cname : str, qobj : optional(dict)=None, *args, **kwargs): + qflags = kwargs.get("qflags", 0) + kwargs["qflags"] = qflags | JBQRYCOUNT + return self.find(cname, qobj, *args, **kwargs) + + def count(self, cname : str, qobj : optional(dict)=None, *args, **kwargs): + qflags = kwargs.get("qflags", 0) + kwargs["qflags"] = qflags | JBQRYCOUNT + return self.find(cname, qobj, *args, **kwargs) + @typecheck + def ensureCollection(self, cname : str, **kwargs): + return self.__ejdb.ensureCollection(cname, **kwargs) + @typecheck + def dropCollection(self, cname : str, **kwargs): + return self.__ejdb.dropCollection(cname, **kwargs) diff --git a/pyejdb/setup.py b/pyejdb/setup.py index bbe9add..5998cc5 100755 --- a/pyejdb/setup.py +++ b/pyejdb/setup.py @@ -1,4 +1,4 @@ -#!/usr/bin/python3.2 +#!/usr/bin/python3 from platform import python_version from distutils.command.build_ext import build_ext as _build_ext @@ -110,7 +110,6 @@ setup( "Intended Audience :: Developers", "License :: OSI Approved :: GNU Lesser General Public License (LGPL)", "Operating System :: POSIX", - "Programming Language :: Python :: 2.7", "Programming Language :: Python :: 3.2", "Topic :: Software Development :: Libraries" ] diff --git a/pyejdb/src/pdbcursor.c b/pyejdb/src/pdbcursor.c index 5605e7a..53e6658 100644 --- a/pyejdb/src/pdbcursor.c +++ b/pyejdb/src/pdbcursor.c @@ -1,5 +1,7 @@ #include "pyejdb.h" +static PyObject* PDBCursor_close(PDBCursor *self); + static PyObject* PDBCursor_tp_new(PyTypeObject *type, PyObject *db, TCLIST *res) { PDBCursor *self = (PDBCursor *) type->tp_alloc(type, 0); if (!self) { @@ -8,7 +10,6 @@ static PyObject* PDBCursor_tp_new(PyTypeObject *type, PyObject *db, TCLIST *res) Py_INCREF(db); self->db = db; self->res = res; - self->pos = 0; return (PyObject *) self; } @@ -23,10 +24,87 @@ static int PDBCursor_tp_clear(PDBCursor *self) { } static void PDBCursor_tp_dealloc(PDBCursor *self) { + PDBCursor_close(self); + Py_TYPE(self)->tp_free((PyObject *) self); +} + +static PyObject* PDBCursor_get(PDBCursor *self, PyObject *args) { + int idx, bsdatasz; + void *bsdata; + if (!PyArg_ParseTuple(args, "i:PDBCursor_get", &idx)) { + return NULL; + } + if (!self->res) { + return set_error(Error, "Cursor closed"); + } + if (idx < 0 || idx >= TCLISTNUM(self->res)) { + return set_error(PyExc_IndexError, "Invalid cursor index"); + } + TCLISTVAL(bsdata, self->res, idx, bsdatasz); + return void_to_bytes(bsdata, bsdatasz); +} + +static PyObject* PDBCursor_length(PDBCursor *self) { + if (!self->res) { + return set_error(Error, "Cursor closed"); + } + return PyLong_FromLong(TCLISTNUM(self->res)); +} + +static PyObject* PDBCursor_close(PDBCursor *self) { if (self->res) { tclistdel(self->res); self->res = NULL; } PDBCursor_tp_clear(self); - Py_TYPE(self)->tp_free((PyObject *) self); + Py_RETURN_NONE; } + +static PyMethodDef PDBCursor_tp_methods[] = { + {"get", (PyCFunction) PDBCursor_get, METH_VARARGS, NULL}, + {"length", (PyCFunction) PDBCursor_length, METH_NOARGS, NULL}, + {"close", (PyCFunction) PDBCursor_close, METH_NOARGS, NULL}, + {NULL} +}; + + +static PyTypeObject DBCursorType = { + PyVarObject_HEAD_INIT(NULL, 0) + "_pyejdb.DBCursor", /*tp_name*/ + sizeof (PDBCursor), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor) PDBCursor_tp_dealloc, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /*tp_flags*/ + 0, /*tp_doc*/ + (traverseproc) PDBCursor_tp_traverse, /*tp_traverse*/ + (inquiry) PDBCursor_tp_clear, /*tp_clear*/ + 0, /*tp_richcompare*/ + 0, /*tp_weaklistoffset*/ + 0, /*tp_iter*/ + 0, /*tp_iternext*/ + PDBCursor_tp_methods, /*tp_methods*/ + 0, /*tp_members*/ + 0, /*tp_getsets*/ + 0, /*tp_base*/ + 0, /*tp_dict*/ + 0, /*tp_descr_get*/ + 0, /*tp_descr_set*/ + 0, /*tp_dictoffset*/ + 0, /*tp_init*/ + 0, /*tp_alloc*/ + 0, /*tp_new*/ +}; diff --git a/pyejdb/src/pejdb.c b/pyejdb/src/pejdb.c index ba1ca0f..cbd3977 100644 --- a/pyejdb/src/pejdb.c +++ b/pyejdb/src/pejdb.c @@ -60,6 +60,42 @@ static PyObject* EJDB_sync(PEJDB *self) { Py_RETURN_NONE; } +static PyObject* EJDB_dropCollection(PEJDB *self, PyObject *args, PyObject *kwargs) { + const char *cname; + PyObject *prune = Py_False; + static char *kwlist[] = {"cname", "prune", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|O:EJDB_dropCollection", kwlist, + &cname, &prune)) { + return NULL; + } + if (!ejdbrmcoll(self->ejdb, cname, (prune == Py_True))) { + return set_ejdb_error(self->ejdb); + } + Py_RETURN_NONE; +} + +static PyObject* EJDB_ensureCollection(PEJDB *self, PyObject *args, PyObject *kwargs) { + const char *cname; + int cachedrecords = 0, records = 0; + PyObject *compressed = Py_False, *large = Py_False; + + static char *kwlist[] = {"cname", "cachedrecords", "compressed", "large", "records", NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|iOOi:EJDB_ensureCollection", kwlist, + &cname, &cachedrecords, &compressed, &large, &records)) { + return NULL; + } + EJCOLLOPTS jcopts = {NULL}; + jcopts.cachedrecords = (cachedrecords > 0) ? cachedrecords : 0; + jcopts.compressed = (compressed == Py_True); + jcopts.large = (large == Py_True); + jcopts.records = (records > 0) ? records : 0; + EJCOLL *coll = ejdbcreatecoll(self->ejdb, cname, &jcopts); + if (!coll) { + return set_ejdb_error(self->ejdb); + } + Py_RETURN_NONE; +} + static PyObject* EJDB_save(PEJDB *self, PyObject *args, PyObject *kwargs) { const char *cname; PyObject *merge = Py_False; @@ -172,6 +208,8 @@ static PyObject* EJDB_find(PEJDB *self, PyObject *args) { int qflags = 0; TCLIST *qres = NULL; EJCOLL *coll = NULL; + PyObject *pyret = NULL; + uint32_t count = 0; //cname, qobj, qobj, orarr, hints, qflags if (!PyArg_ParseTuple(args, "sOOOi:EJDB_find", @@ -231,11 +269,27 @@ static PyObject* EJDB_find(PEJDB *self, PyObject *args) { } } } + if (!coll) { //No collection -> no results - goto finish; + qres = (qflags & JBQRYCOUNT) ? NULL : tclistnew2(1); //empty results + } else { + Py_BEGIN_ALLOW_THREADS + qres = ejdbqryexecute(coll, q, &count, qflags, NULL); + Py_END_ALLOW_THREADS + if (ejdbecode(self->ejdb) != TCESUCCESS) { + err = true; + set_ejdb_error(self->ejdb); + goto finish; + } } - //TODO + if (qres) { + pyret = PDBCursor_tp_new(&DBCursorType, (PyObject *) self, qres); + if (!pyret) { + err = true; + goto finish; + } + } finish: for (Py_ssize_t i = 0; i < orsz; ++i) { @@ -250,10 +304,10 @@ finish: ejdbquerydel(q); } if (err) { + Py_XDECREF(pyret); return NULL; } else { - //cursor? - Py_RETURN_NONE; + return (pyret) ? pyret : PyLong_FromUnsignedLong(count); } } @@ -267,6 +321,8 @@ static PyMethodDef EJDB_tp_methods[] = { {"remove", (PyCFunction) EJDB_remove, METH_VARARGS, NULL}, {"sync", (PyCFunction) EJDB_sync, METH_NOARGS, NULL}, {"find", (PyCFunction) EJDB_find, METH_VARARGS, NULL}, + {"ensureCollection", (PyCFunction) EJDB_ensureCollection, METH_VARARGS | METH_KEYWORDS, NULL}, + {"dropCollection", (PyCFunction) EJDB_dropCollection, METH_VARARGS | METH_KEYWORDS, NULL}, {NULL} }; diff --git a/pyejdb/src/pyejdb.c b/pyejdb/src/pyejdb.c index c79f0cd..f23bdba 100644 --- a/pyejdb/src/pyejdb.c +++ b/pyejdb/src/pyejdb.c @@ -11,14 +11,12 @@ typedef struct { PyObject_HEAD PyObject *db; TCLIST *res; - int pos; } PDBCursor; #include "pdbcursor.c" #include "pejdb.c" - PyDoc_STRVAR(ejdb_m_doc, "EJDB http://ejdb.org"); PyDoc_STRVAR(ejdb_version_doc, "version() -> str\n\nReturns the version string of the underlying EJDB library."); @@ -26,29 +24,30 @@ static PyObject* ejdb_version(PyObject *module) { return PyUnicode_FromString(tcversion); } -/* cabinet_module.m_methods */ +/* pyejdb.m_methods */ static PyMethodDef pyejdb_m_methods[] = { - {"version", (PyCFunction) ejdb_version, METH_NOARGS, ejdb_version_doc}, + {"ejdb_version", (PyCFunction) ejdb_version, METH_NOARGS, ejdb_version_doc}, {NULL} /* Sentinel */ }; /* pyejdb_module */ static PyModuleDef pyejdb_module = { PyModuleDef_HEAD_INIT, - "_pyejdb", /*m_name*/ - ejdb_m_doc, /*m_doc*/ - -1, /*m_size*/ - pyejdb_m_methods, /*m_methods*/ + "_pyejdb", /*m_name*/ + ejdb_m_doc, /*m_doc*/ + -1, /*m_size*/ + pyejdb_m_methods, /*m_methods*/ }; PyObject* init_pyejdb(void) { PyObject *pyejdb; - if (PyType_Ready(&EJDBType)) { + if (PyType_Ready(&EJDBType) || + PyType_Ready(&DBCursorType) + ) { return NULL; } pyejdb = PyModule_Create(&pyejdb_module); - if (!pyejdb) { return NULL; } @@ -62,6 +61,7 @@ PyObject* init_pyejdb(void) { /* adding types and constants */ if ( PyModule_AddType(pyejdb, "EJDB", &EJDBType) || + PyModule_AddType(pyejdb, "DBCursor", &DBCursorType) || PyModule_AddIntMacro(pyejdb, JBOREADER) || PyModule_AddIntMacro(pyejdb, JBOWRITER) || diff --git a/pyejdb/src/pyejdb.h b/pyejdb/src/pyejdb.h index e3617a5..da91548 100644 --- a/pyejdb/src/pyejdb.h +++ b/pyejdb/src/pyejdb.h @@ -84,6 +84,10 @@ int bytes_to_void(PyObject *pyvalue, void **value, int *value_len) { return 0; } +PyObject* void_to_bytes(const void *value, int value_size) { + return PyBytes_FromStringAndSize((char *) value, (Py_ssize_t) value_size); +} + /* error handling utils */ PyObject* set_error(PyObject *type, const char *message) { PyErr_SetString(type, message); diff --git a/pyejdb/test/test_one.py b/pyejdb/test/test_one.py index 5b1240c..5ca33d2 100644 --- a/pyejdb/test/test_one.py +++ b/pyejdb/test/test_one.py @@ -9,6 +9,8 @@ class TestOne(unittest.TestCase): @classmethod def setUpClass(cls): + print("pyejdb version: %s" % pyejdb.version) + print("libejdb_version: %s" % pyejdb.libejdb_version) cls._ejdb = pyejdb.EJDB("testdb", pyejdb.DEFAULT_OPEN_MODE | pyejdb.JBOTRUNC) def test(self): @@ -22,13 +24,34 @@ class TestOne(unittest.TestCase): self.assertEqual(doc["_id"], ldoc["_id"]) self.assertEqual(doc["foo"], ldoc["foo"]) self.assertEqual(doc["foo2"], ldoc["foo2"]) + + cur = ejdb.find("foocoll", {"foo": "bar"}, hints={"$fields": {"foo2": 0}}) + self.assertEqual(len(cur), 1) + d = cur[0] + self.assertTrue(d is not None) + self.assertEqual(d["_id"], ldoc["_id"]) + + with ejdb.find("foocoll") as cur2: + d = cur2[0] + self.assertTrue(d is not None) + self.assertEqual(d["_id"], ldoc["_id"]) + + self.assertEqual(ejdb.findOne("foocoll")["foo"], "bar") + self.assertTrue(ejdb.findOne("foocoll2") is None) + + self.assertEqual(ejdb.count("foocoll"), 1) + self.assertEqual(ejdb.count("foocoll2"), 0) + ejdb.remove("foocoll", doc["_id"]) ldoc = ejdb.load("foocoll", doc["_id"]) self.assertTrue(ldoc is None) ejdb.sync() + ejdb.ensureCollection("ecoll1", records=90000, large=False) + ejdb.dropCollection("ecoll1", prune=True) - ejdb.find("foocoll", {}, {}, hints={}) + def test2(self): + pass @classmethod def tearDownClass(cls): -- 2.7.4