mygpo/couch.py

   1 from operator import itemgetter
   2 from collections import namedtuple
   3
   4 from django.conf import settings
   5
   6 from couchdbkit import *
   7
   8 class BulkException(Exception):
   9
  10     def __init__(self, errors):
  11         self.errors = errors
  12
  13
  14 BulkError = namedtuple('BulkError', 'doc error reason')
  15
  16
  17 def __default_reload(db, obj):
  18     _id = obj._id
  19
  20     if isinstance(obj, Document):
  21         return obj.__class__.get(_id)
  22     else:
  23         return db[_id]
  24
  25
  26 __get_obj = itemgetter(0)
  27
  28 def bulk_save_retry(db, obj_funs, reload_f=__default_reload):
  29     """ Saves multiple documents and retries failed ones
  30
  31     Objects to be saved are passed as (obj, mod_f), where obj is the CouchDB
  32     and mod_f is the modification function that should be applied to it.
  33
  34     If saving a document fails, it is again fetched from the database, the
  35     modification function is applied again and saving is retried. """
  36
  37     errors = []
  38
  39     while True:
  40
  41         # apply modification function (and keep funs)
  42         obj_funs = [(f(o), f) for (o, f) in obj_funs]
  43
  44         # filter those with obj None
  45         obj_funs = filter(lambda of: __get_obj(of) is not None, obj_funs)
  46
  47         # extract objects
  48         objs = map(__get_obj, obj_funs)
  49
  50         if not objs:
  51             return
  52
  53         try:
  54             db.save_docs(objs)
  55             return
  56
  57         except BulkSaveError as ex:
  58
  59             new_obj_funs = []
  60             for res, (obj, f) in zip(ex.results, obj_funs):
  61                 if res.get('error', False) == 'conflict':
  62
  63                     # reload conflicted object
  64                     obj = reload_f(db, obj)
  65                     new_obj_funs.append( (obj, f) )
  66
  67                 elif res.get('error', False):
  68                     # don't retry other errors
  69                     err = BulkError(obj, res['error'], res.get('reason', None))
  70                     errors.append(err)
  71
  72             obj_funs = new_obj_funs
  73
  74     if errors:
  75         raise BulkException(errors)
  76
  77
  78
  79 def get_main_database():
  80     db_url = settings.COUCHDB_DATABASES[0][1]
  81     return Database(db_url)