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
   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 from mygpo.db.couchdb.episode_state import episode_states_count, \
  13          get_nth_episode_state, get_duplicate_episode_states
  14
  15
  16 class Command(BaseCommand):
  17     """ Merge duplicate EpisodeUserState documents """
  18
  19     option_list = BaseCommand.option_list + (
  20         make_option('--skip', action='store', type=int, dest='skip', default=0,
  21            help="Number of states to skip"),
  22     )
  23
  24     def handle(self, *args, **options):
  25
  26         skip = options.get('skip')
  27         total = episode_states_count()
  28         db = get_main_database()
  29
  30         actions = Counter()
  31         actions['merged'] = 0
  32
  33
  34         for n in count(skip):
  35
  36             first = get_nth_episode_state(n)
  37             states = get_duplicate_episode_states(first.user, first.episode)
  38
  39             l1 = len(states)
  40             # we don't want to delete this one
  41             states.remove(first)
  42
  43             assert len(states) == l1-1
  44
  45             if states:
  46                 updater = get_updater(states)
  47
  48                 obj_funs = [(first, updater)] + [(state, do_delete) for state in states]
  49
  50                 bulk_save_retry(db, obj_funs)
  51
  52                 merged = len(states)-1
  53                 actions['merged'] += merged
  54                 total -= merged
  55
  56             status_str = ', '.join('%s: %d' % x for x in actions.items())
  57             progress(n+1, total, status_str)
  58
  59
  60 def get_updater(states):
  61
  62     actions = set()
  63     settings = dict()
  64     merged_ids = set()
  65     chapters = set()
  66
  67     for state in states:
  68         actions.union(set(state.actions))
  69         settings.update(state.settings)
  70         merged_ids.union(set(state.merged_ids + [state._id]))
  71         chapters.union(set(state.chapters))
  72
  73     return partial(do_update, list(actions), settings, list(merged_ids), list(chapters))
  74
  75
  76 def do_update(actions, settings, merged_ids, chapters, state):
  77     state.add_actions(actions)
  78     # overwrite settings in old_state with state's settings
  79     state.settings = settings.update(state.settings or {})
  80     state.merged_ids = list(set(state.merged_ids + merged_ids))
  81     state.chapters = list(set(state.chapters + chapters))
  82     return state
  83
  84
  85 def do_delete(state):
  86     # remove all attributes
  87     for attr in filter(lambda n: not n.startswith('_'), dir(state)):
  88         try:
  89             delattr(state, attr)
  90         except AttributeError:
  91             pass
  92
  93     state._deleted = True
  94     return state