1# Python interface to CouchStore library
2
3import errno
4import inspect
5import os
6import sys
7import traceback
8
9try:
10    import ctypes
11except ImportError:
12    cb_path = '/opt/couchbase/lib/python'
13    while cb_path in sys.path:
14        sys.path.remove(cb_path)
15    try:
16        import ctypes
17    except ImportError:
18        sys.exit('error: could not import ctypes module')
19    else:
20        sys.path.insert(0, cb_path)
21
22# Load the couchstore library and customize return types:
23_lib_dir = [
24    # cbbackup / cbrestore
25    os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))[:-6],
26    # couchstore python tests - uses the LD_LIBRARY_PATH set
27    # via CMakeLists.txt to explicitly identify the path on OSX
28    os.environ.get('LD_LIBRARY_PATH', '.')
29    ]
30osx_path = [os.path.join(path, 'libcouchstore.dylib') for path in _lib_dir]
31
32for lib in ['libcouchstore.so',      # Linux
33            'couchstore.dll',        # Windows
34            'libcouchstore-1.dll'    # Windows (pre-CMake)
35            ] + osx_path:            # Mac OS
36    try:
37        _lib = ctypes.CDLL(lib)
38        break
39    except OSError, err:
40        continue
41else:
42    raise ImportError("Failed to locate suitable couchstore shared library")
43
44
45_lib.couchstore_strerror.restype = ctypes.c_char_p
46
47
48class CouchStoreException(Exception):
49    """Exceptions raised by CouchStore APIs."""
50    def __init__(self, errcode):
51        Exception.__init__(self, _lib.couchstore_strerror(errcode))
52        self.code = errcode
53
54
55### INTERNAL FUNCTIONS:
56
57def _check(err):
58    if err == 0:
59        return
60    elif err == -3:
61        raise MemoryError()
62    elif err == -5:
63        raise KeyError()
64    elif err == -11:
65        raise OSError(errno.ENOENT)
66    else:
67        raise CouchStoreException(err)
68
69
70def _toString(key):
71    if not isinstance(key, basestring):
72        raise TypeError(key)
73    return str(key)
74
75
76### INTERNAL STRUCTS:
77
78class SizedBuf(ctypes.Structure):
79    _fields_ = [("buf", ctypes.POINTER(ctypes.c_char)), ("size", ctypes.c_size_t)]
80
81    def __init__(self, string):
82        if string is not None:
83            string = _toString(string)
84            length = len(string)
85            buf = ctypes.create_string_buffer(string, length)
86            ctypes.Structure.__init__(self, buf, length)
87        else:
88            ctypes.Structure.__init__(self, None, 0)
89
90    def __str__(self):
91        return ctypes.string_at(self.buf, self.size)
92
93
94class DocStruct(ctypes.Structure):
95    _fields_ = [("id", SizedBuf), ("data", SizedBuf)]
96
97
98class DocInfoStruct(ctypes.Structure):
99    _fields_ = [("id", SizedBuf),
100                ("db_seq", ctypes.c_ulonglong),
101                ("rev_seq", ctypes.c_ulonglong),
102                ("rev_meta", SizedBuf),
103                ("deleted", ctypes.c_int),
104                ("content_meta", ctypes.c_ubyte),
105                ("bp", ctypes.c_ulonglong),
106                ("size", ctypes.c_size_t)]
107
108
109class LocalDocStruct(ctypes.Structure):
110    _fields_ = [("id", SizedBuf),
111                ("json", SizedBuf),
112                ("deleted", ctypes.c_int)]
113
114
115class DbInfoStruct(ctypes.Structure):
116    _fields_ = [("filename", ctypes.c_char_p),
117                ("last_sequence", ctypes.c_ulonglong),
118                ("doc_count", ctypes.c_ulonglong),
119                ("deleted_count", ctypes.c_ulonglong),
120                ("space_used", ctypes.c_ulonglong),
121                ("header_position", ctypes.c_ulonglong)]
122
123
124class CounterStruct(ctypes.Structure):
125    _fields_ = [("count", ctypes.c_ulonglong)]
126
127
128class DocumentInfo(object):
129    """Metadata of a document in a CouchStore database."""
130
131    # Values for contentType:
132    IS_JSON = 0
133    INVALID_JSON = 1
134    INVALID_JSON_KEY = 2
135    NON_JSON = 3
136
137    def __init__(self, id):
138        self.id = id
139        self.deleted = False
140        self.contentType = DocumentInfo.NON_JSON
141        self.revSequence = 0
142
143    @staticmethod
144    def _fromStruct(info, store=None):
145        self = DocumentInfo(str(info.id))
146        self.store = store
147        self.sequence = info.db_seq
148        self.revSequence = info.rev_seq
149        self.revMeta = str(info.rev_meta)
150        self.deleted = (info.deleted != 0)
151        self.contentType = info.content_meta & 0x0F
152        self.compressed = (info.content_meta & 0x80) != 0
153        self._bp = info.bp
154        self.physSize = info.size
155        return self
156
157    def _asStruct(self):
158        struct = DocInfoStruct(SizedBuf(self.id))
159        if hasattr(self, "sequence"):
160            struct.db_seq = self.sequence
161        if hasattr(self, "revMeta"):
162            struct.rev_meta = SizedBuf(self.revMeta)
163        struct.rev_seq = self.revSequence
164        struct.deleted = self.deleted
165        struct.content_meta = self.contentType & 0x0F
166        if hasattr(self, "compressed") and self.compressed:
167            struct.content_meta |= 0x80
168        if hasattr(self, "_bp"):
169            struct.bp = self._bp
170        if hasattr(self, "physSize"):
171            struct.size = self.physSize
172        return struct
173
174    def __str__(self):
175        return "DocumentInfo('%s', %d bytes)" % (self.id, self.physSize)
176
177    def __repr__(self):
178        return "DocumentInfo('%s', %d bytes)" % (self.id, self.physSize)
179
180    def dump(self):
181        return "DocumentInfo('%s', %d bytes, seq=%d, revSeq=%d, deleted=%s, " \
182               "contentType=%d, compressed=%d, bp=%d)" % \
183            (self.id, self.physSize, self.sequence, self.revSequence,
184             self.deleted, self.contentType, self.compressed, self._bp)
185
186    def getContents(self, options=0):
187        """Fetches and returns the contents of a DocumentInfo returned from
188        CouchStore's getInfo or getInfoBySequence methods."""
189        if not hasattr(self, "store") or not hasattr(self, "_bp"):
190            raise Exception("Contents unknown")
191        info = self._asStruct()
192        docptr = ctypes.pointer(DocStruct())
193        _check(_lib.couchstore_open_doc_with_docinfo(self.store,
194                                                     ctypes.byref(info),
195                                                     ctypes.byref(docptr),
196                                                     ctypes.c_uint64(options)))
197        contents = str(docptr.contents.data)
198        _lib.couchstore_free_document(docptr)
199        return contents
200
201
202class LocalDocs(object):
203    """Collection that represents the local documents of a CouchStore."""
204
205    def __init__(self, couchstore):
206        self.couchstore = couchstore
207
208    def __getitem__(self, key):
209        """Returns the contents of a local document (as a string) given its ID.
210        """
211        id = _toString(key)
212        docptr = ctypes.pointer(LocalDocStruct())
213        err = _lib.couchstore_open_local_document(self.couchstore,
214                                                  id,
215                                                  ctypes.c_size_t(len(id)),
216                                                  ctypes.byref(docptr))
217        if err == -5 or (err == 0 and docptr.contents.deleted):
218            raise KeyError(id)
219        _check(err)
220        value = str(docptr.contents.json)
221        _lib.couchstore_free_document(docptr)
222        return value
223
224    def __setitem__(self, key, value):
225        """Saves a local document with the given ID, or deletes it if the value
226        is None."""
227        idbuf = SizedBuf(key)
228        doc = LocalDocStruct(idbuf)
229        if value is not None:
230            doc.json = SizedBuf(value)
231        else:
232            doc.deleted = True
233        _check(_lib.couchstore_save_local_document(self.couchstore,
234                                                   ctypes.byref(doc)))
235
236    def __delitem__(self, key):
237        self.__setitem__(key, None)
238
239
240class CouchStore(object):
241    """Interface to a CouchStore database."""
242
243    def __init__(self, path, mode=None, unbuffered=False):
244        """Creates a CouchStore at a given path. The option mode parameter can
245        be 'r' for read-only access, or 'c' to create the file if it doesn't
246        already exist."""
247        if mode == 'r':
248            flags = 2  # RDONLY
249        elif mode == 'c':
250            flags = 1  # CREATE
251        else:
252            flags = 0
253
254        if unbuffered:
255            flags |= 8 ## UNBUFFERED
256
257        db = ctypes.c_void_p()
258        _check(_lib.couchstore_open_db(path,
259                                       ctypes.c_uint64(flags),
260                                       ctypes.byref(db)))
261        self._as_parameter_ = db
262        self.path = path
263
264    def __del__(self):
265        self.close()
266
267    def close(self):
268        """Closes the CouchStore."""
269        if hasattr(self, "_as_parameter_"):
270            _lib.couchstore_close_file(self)
271            _lib.couchstore_free_db(self)
272            del self._as_parameter_
273
274    def __str__(self):
275        return "CouchStore(%s)" % self.path
276
277    def getDbInfo(self):
278        """Returns an object with information about the database. Its
279        properties are filename, last_sequence, doc_count, deleted_count,
280        space_used, header_position."""
281        info = DbInfoStruct()
282        _check(_lib.couchstore_db_info(self, ctypes.byref(info)))
283        return info
284
285    def rewindHeader(self):
286        """Rewinds the database handle to the next-oldest committed header.
287        Closes the handle if none can be found"""
288        if hasattr(self, "_as_parameter_"):
289            err = _lib.couchstore_rewind_db_header(self)
290            if err != 0:
291                del self._as_parameter_
292            _check(err)
293
294    COMPRESS = 1
295
296    def save(self, id, data, options=0):
297        """Saves a document with the given ID. Returns the sequence number."""
298        if isinstance(id, DocumentInfo):
299            infoStruct = id._asStruct()
300            idbuf = infoStruct.id
301        else:
302            idbuf = SizedBuf(id)
303            infoStruct = DocInfoStruct(idbuf)
304        if data is not None:
305            doc = DocStruct(idbuf, SizedBuf(data))
306            docref = ctypes.byref(doc)
307            if options & CouchStore.COMPRESS:
308                infoStruct.content_meta |= 0x80
309        else:
310            docref = None
311        _check(_lib.couchstore_save_document(self, docref,
312                                             ctypes.byref(infoStruct),
313                                             ctypes.c_uint64(options)))
314        if isinstance(id, DocumentInfo):
315            id.sequence = infoStruct.db_seq
316        return infoStruct.db_seq
317
318    def saveMultiple(self, ids, datas, options=0):
319        """Saves multiple documents. 'ids' is an array of either strings or
320        DocumentInfo objects. 'datas' is a parallel array of value strings (or
321        None, in which case the documents will be deleted.) Returns an array of
322        new sequence numbers."""
323        n = len(ids)
324        docStructs = (ctypes.POINTER(DocStruct) * n)()
325        infoStructs = (ctypes.POINTER(DocInfoStruct) * n)()
326        for i in xrange(0, n):
327            id = ids[i]
328            if isinstance(id, DocumentInfo):
329                info = id._asStruct()
330            else:
331                info = DocInfoStruct(SizedBuf(id))
332            doc = DocStruct(info.id)
333            if datas and datas[i]:
334                doc.data = SizedBuf(datas[i])
335            else:
336                info.deleted = True
337            infoStructs[i] = ctypes.pointer(info)
338            docStructs[i] = ctypes.pointer(doc)
339        _check(_lib.couchstore_save_documents(self,
340                                              ctypes.byref(docStructs),
341                                              ctypes.byref(infoStructs),
342                                              ctypes.c_uint(n),
343                                              ctypes.c_uint64(options)))
344        return [info.contents.db_seq for info in infoStructs]
345    pass
346
347    def commit(self):
348        """Ensures all saved data is flushed to disk."""
349        _check(_lib.couchstore_commit(self))
350
351    DECOMPRESS = 1
352
353    def get(self, id, options=0):
354        """Returns the contents of a document (as a string) given its ID."""
355        id = _toString(id)
356        docptr = ctypes.pointer(DocStruct())
357        err = _lib.couchstore_open_document(self,
358                                            id,
359                                            ctypes.c_size_t(len(id)),
360                                            ctypes.byref(docptr),
361                                            options)
362        if err == -5:
363            raise KeyError(id)
364        _check(err)
365        data = str(docptr.contents.data)
366        _lib.couchstore_free_document(docptr)
367        return data
368
369    def __getitem__(self, key):
370        return self.get(key)
371
372    def __setitem__(self, key, value):
373        self.save(key, value)
374
375    def __delitem__(self, key):
376        self.save(key, None)
377
378    # Getting document info:
379
380    def _infoPtrToDoc(self, key, infoptr, err):
381        if err == -5:
382            raise KeyError(key)
383        _check(err)
384        info = infoptr.contents
385        if info is None:
386            return None
387        doc = DocumentInfo._fromStruct(info, self)
388        _lib.couchstore_free_docinfo(infoptr)
389        return doc
390
391    def getInfo(self, id):
392        """Returns the DocumentInfo object with the given ID."""
393        id = _toString(id)
394        infoptr = ctypes.pointer(DocInfoStruct())
395        err = _lib.couchstore_docinfo_by_id(self,
396                                            id,
397                                            ctypes.c_size_t(len(id)),
398                                            ctypes.byref(infoptr))
399        return self._infoPtrToDoc(id, infoptr, err)
400
401    def getInfoBySequence(self, sequence):
402        """Returns the DocumentInfo object with the given sequence number."""
403        infoptr = ctypes.pointer(DocInfoStruct())
404        err = _lib.couchstore_docinfo_by_sequence(self,
405                                                  ctypes.c_ulonglong(sequence),
406                                                  ctypes.byref(infoptr))
407        return self._infoPtrToDoc(sequence, infoptr, err)
408
409    # Iterating:
410
411    ITERATORFUNC = ctypes.CFUNCTYPE(ctypes.c_int,
412                                    ctypes.c_void_p,
413                                    ctypes.POINTER(DocInfoStruct),
414                                    ctypes.c_void_p)
415
416    def forEachChange(self, since, fn):
417        """Calls the function "fn" once for every document sequence since the
418        "since" parameter. The single parameter to "fn" will be a DocumentInfo
419        object. You can call getContents() on it to get the document contents.
420        """
421        def callback(dbPtr, docInfoPtr, context):
422            fn(DocumentInfo._fromStruct(docInfoPtr.contents, self))
423            return 0
424        _check(_lib.couchstore_changes_since(self,
425                                             ctypes.c_uint64(since),
426                                             ctypes.c_uint64(0),
427                                             CouchStore.ITERATORFUNC(callback),
428                                             ctypes.c_void_p(0)))
429
430    def changesSince(self, since):
431        """Returns an array of DocumentInfo objects, for every document that's
432        changed since the sequence number "since"."""
433        changes = []
434        self.forEachChange(since, lambda docInfo: changes.append(docInfo))
435        return changes
436
437    def forEachDoc(self, startKey, endKey, fn):
438        def callback(dbPtr, docInfoPtr, context):
439            fn(DocumentInfo._fromStruct(docInfoPtr.contents, self))
440            return 0
441
442        ids = (SizedBuf * 2)()
443        numIDs = 1
444        if startKey:
445            ids[0] = SizedBuf(startKey)
446        if endKey:
447            ids[1] = SizedBuf(endKey)
448            numIDs = 2
449        _check(_lib.couchstore_docinfos_by_id(self,
450                                              ids,
451                                              ctypes.c_uint(numIDs),
452                                              ctypes.c_uint64(1),
453                                              CouchStore.ITERATORFUNC(callback),
454                                              ctypes.c_void_p(0)))
455
456    def changesCount(self, minimum, maximum):
457        cstruct = CounterStruct()
458        err = _lib.couchstore_changes_count(self,
459                                            ctypes.c_uint64(minimum),
460                                            ctypes.c_uint64(maximum),
461                                            ctypes.pointer(cstruct))
462        _check(err)
463        return cstruct.count
464
465    @property
466    def localDocs(self):
467        """A simple dictionary-like object that accesses the CouchStore's local
468        documents."""
469        return LocalDocs(self)
470