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 from collections import Counter
   5
   6 from django.core.management.base import BaseCommand
   7
   8 from mygpo.utils import progress
   9 from mygpo.db.couchdb import bulk_save_retry
  10 from mygpo.db.couchdb.episode_state import episode_states_count, \
  11          get_nth_episode_state, get_duplicate_episode_states
  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 = episode_states_count()
  26
  27         actions = Counter()
  28         actions['merged'] = 0
  29
  30
  31         for n in count(skip):
  32
  33             first = get_nth_episode_state(n)
  34             if first is None:
  35                 break
  36             states = get_duplicate_episode_states(first.user, first.episode)
  37
  38             l1 = len(states)
  39             # we don't want to delete this one
  40             states.remove(first)
  41
  42             assert len(states) == l1-1
  43
  44             if states:
  45                 updater = get_updater(states)
  46
  47                 obj_funs = [(first, updater)] + [(state, do_delete) for state in states]
  48
  49                 bulk_save_retry(obj_funs)
  50
  51                 merged = len(states)-1
  52                 actions['merged'] += merged
  53                 total -= merged
  54
  55             status_str = ', '.join('%s: %d' % x for x in actions.items())
  56             progress(n+1, total, status_str)
  57
  58
  59 def get_updater(states):
  60
  61     actions = set()
  62     settings = dict()
  63     merged_ids = set()
  64     chapters = set()
  65
  66     for state in states:
  67         actions.union(set(state.actions))
  68         settings.update(state.settings)
  69         merged_ids.union(set(state.merged_ids + [state._id]))
  70         chapters.union(set(state.chapters))
  71
  72     return partial(do_update, list(actions), settings, list(merged_ids), list(chapters))
  73
  74
  75 def do_update(actions, settings, merged_ids, chapters, state):
  76     state.add_actions(actions)
  77     # overwrite settings in old_state with state's settings
  78     state.settings = settings.update(state.settings or {})
  79     state.merged_ids = list(set(state.merged_ids + merged_ids))
  80     state.chapters = list(set(state.chapters + chapters))
  81     return state
  82
  83
  84 def do_delete(state):
  85     # remove all attributes
  86     for attr in filter(lambda n: not n.startswith('_'), dir(state)):
  87         try:
  88             delattr(state, attr)
  89         except AttributeError:
  90             pass
  91
  92     state._deleted = True
  93     return state