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(obj_funs, db=None, 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     db = db or get_main_database()
  38     errors = []
  39
  40     while True:
  41
  42         # apply modification function (and keep funs)
  43         obj_funs = [(f(o), f) for (o, f) in obj_funs]
  44
  45         # filter those with obj None
  46         obj_funs = filter(lambda of: __get_obj(of) is not None, obj_funs)
  47
  48         # extract objects
  49         objs = map(__get_obj, obj_funs)
  50
  51         if not objs:
  52             return
  53
  54         try:
  55             db.save_docs(objs)
  56             return
  57
  58         except BulkSaveError as ex:
  59
  60             new_obj_funs = []
  61             for res, (obj, f) in zip(ex.results, obj_funs):
  62                 if res.get('error', False) == 'conflict':
  63
  64                     # reload conflicted object
  65                     obj = reload_f(db, obj)
  66                     new_obj_funs.append( (obj, f) )
  67
  68                 elif res.get('error', False):
  69                     # don't retry other errors
  70                     err = BulkError(obj, res['error'], res.get('reason', None))
  71                     errors.append(err)
  72
  73             obj_funs = new_obj_funs
  74
  75     if errors:
  76         raise BulkException(errors)
  77
  78
  79
  80 def get_main_database():
  81     db_url = settings.COUCHDB_DATABASES[0][1]
  82     return Database(db_url)