mygpo/maintenance/management/commands/merge-episode-states.py

   1 from optparse import make_option
   2 from itertools import count
   3 from functools import partial
   4
   5 from django.core.management.base import BaseCommand
   6
   7 from mygpo.utils import progress, multi_request_view
   8 from mygpo.users.models import EpisodeUserState
   9 from mygpo.counter import Counter
  10 from mygpo.maintenance.merge import merge_episode_states
  11 from mygpo.couch import bulk_save_retry, get_main_database
  12
  13
  14 class Command(BaseCommand):
  15     """ Merge duplicate EpisodeUserState documents """
  16
  17     option_list = BaseCommand.option_list + (
  18         make_option('--skip', action='store', type=int, dest='skip', default=0,
  19            help="Number of states to skip"),
  20     )
  21
  22     def handle(self, *args, **options):
  23
  24         skip = options.get('skip')
  25         total = EpisodeUserState.view('episode_states/by_user_episode',
  26                 limit=0,
  27             ).total_rows
  28         db = get_main_database()
  29
  30         actions = Counter()
  31         actions['merged'] = 0
  32
  33
  34         for n in count(skip):
  35
  36             first = EpisodeUserState.view('episode_states/by_user_episode',
  37                     skip         = n,
  38                     include_docs = True,
  39                     limit        = 1,
  40                 )
  41             first = list(first)
  42             if not first:
  43                 break
  44
  45             first = first[0]
  46
  47
  48             states = EpisodeUserState.view('episode_states/by_user_episode',
  49                     key          = [first.user, first.episode],
  50                     include_docs = True,
  51                 )
  52             states = list(states)
  53
  54             l1 = len(states)
  55             # we don't want to delete this one
  56             states.remove(first)
  57
  58             assert len(states) == l1-1
  59
  60             if states:
  61                 updater = get_updater(states)
  62
  63                 obj_funs = [(first, updater)] + [(state, do_delete) for state in states]
  64
  65                 bulk_save_retry(db, obj_funs)
  66
  67                 merged = len(states)-1
  68                 actions['merged'] += merged
  69                 total -= merged
  70
  71             status_str = ', '.join('%s: %d' % x for x in actions.items())
  72             progress(n+1, total, status_str)
  73
  74
  75 def get_updater(states):
  76
  77     actions = set()
  78     settings = dict()
  79     merged_ids = set()
  80     chapters = set()
  81
  82     for state in states:
  83         actions.union(set(state.actions))
  84         settings.update(state.settings)
  85         merged_ids.union(set(state.merged_ids + [state._id]))
  86         chapters.union(set(state.chapters))
  87
  88     return partial(do_update, list(actions), settings, list(merged_ids), list(chapters))
  89
  90
  91 def do_update(actions, settings, merged_ids, chapters, state):
  92     state.add_actions(actions)
  93     # overwrite settings in old_state with state's settings
  94     state.settings = settings.update(state.settings or {})
  95     state.merged_ids = list(set(state.merged_ids + merged_ids))
  96     state.chapters = list(set(state.chapters + chapters))
  97     return state
  98
  99
 100 def do_delete(state):
 101     # remove all attributes
 102     for attr in filter(lambda n: not n.startswith('_'), dir(state)):
 103         try:
 104             delattr(state, attr)
 105         except AttributeError:
 106             pass
 107
 108     state._deleted = True
 109     return state