mygpo/maintenance/management/commands/assign-episode-slugs.py

   1 from datetime import datetime
   2 from optparse import make_option
   3 from collections import Counter
   4
   5 from django.core.management.base import BaseCommand
   6
   7 from couchdbkit.exceptions import ResourceNotFound
   8 from couchdbkit import Consumer
   9
  10 from mygpo.core.models import Podcast, PodcastGroup, Episode
  11 from mygpo.core.slugs import EpisodeSlug, EpisodesMissingSlugs
  12 from mygpo.decorators import repeat_on_conflict
  13 from mygpo.utils import progress
  14 from mygpo.db.couchdb import get_main_database
  15 from mygpo.maintenance.models import CommandStatus, CommandRunStatus
  16 from mygpo.db.couchdb.podcast import podcast_by_id
  17
  18
  19 class Command(BaseCommand):
  20
  21     option_list = BaseCommand.option_list + (
  22         make_option('--since', action='store', type=int, dest='since',
  23             default=0, help="Where to start the operation"),
  24
  25         make_option('--continue', action='store_true', dest='continue',
  26             default=False, help="Continue from last sequence number"),
  27
  28         make_option('--silent', action='store_true', dest='silent',
  29             default=False, help="Don't display any progress output"),
  30     )
  31
  32
  33     def handle(self, *args, **options):
  34
  35         db = get_main_database()
  36         status = self.get_cmd_status()
  37         since = self.get_since(status, options)
  38         objects = self.get_objects(db, since)
  39         actions = Counter()
  40
  41
  42         # create unfinished command run status
  43         run_status = CommandRunStatus()
  44         run_status.timestamp_started = datetime.utcnow()
  45         run_status.start_seq = since
  46         # add it to existing one (if any)
  47         status.runs.append(run_status)
  48         status.save()
  49
  50         total = db.info()['update_seq']
  51
  52         has_slug = lambda x: bool(x.slug)
  53
  54         for seq, obj in objects:
  55             total = db.info()['update_seq']
  56
  57             if isinstance(obj, PodcastGroup):
  58                 podcasts = filter(has_slug, obj.podcasts)
  59
  60             if isinstance(obj, Podcast):
  61                 podcasts = filter(has_slug, [obj])
  62
  63             elif isinstance(obj, Episode):
  64                 if has_slug(obj):
  65                     continue
  66
  67                 podcast = podcast_by_id(obj.podcast)
  68                 if not podcast:
  69                     continue
  70                 podcasts = filter(has_slug, [podcast])
  71
  72             updated = self.handle_podcasts(podcasts)
  73             actions['updated'] += updated
  74
  75             if not options['silent']:
  76                 status_str = ', '.join('%s: %d' % x for x in actions.items())
  77                 progress(seq, total, status_str)
  78
  79
  80         # finish command run status
  81         run_status.timestamp_finished = datetime.utcnow()
  82         run_status.end_seq = total
  83         run_status.status_counter = dict(actions)
  84         # and overwrite existing one (we could keep a longer log here)
  85         status.runs = [run_status]
  86         status.save()
  87
  88
  89
  90     def handle_podcasts(self, podcasts):
  91
  92         updated = 0
  93         for podcast in podcasts:
  94             common_title = podcast.get_common_episode_title()
  95             episodes = EpisodesMissingSlugs(podcast.get_id())
  96
  97             for episode in episodes:
  98                 slug = EpisodeSlug(episode, common_title).get_slug()
  99                 if slug:
 100                     updated += 1
 101                     self.update_obj(obj=episode, slug=slug)
 102
 103         return updated
 104
 105
 106     @repeat_on_conflict(['obj'])
 107     def update_obj(self, obj, slug):
 108         obj.set_slug(slug)
 109         obj.save()
 110
 111
 112     @staticmethod
 113     def get_objects(db, since=0, limit=1, timeout=1000):
 114         consumer = Consumer(db)
 115
 116         while True:
 117             resp = consumer.wait_once(since=since, limit=limit, timeout=timeout,
 118                    include_docs=True, filter='slugs/slug_objects')
 119             results = resp['results']
 120
 121             if not results:
 122                 break
 123
 124             for res in results:
 125                 cls = (PodcastGroup, Podcast, Episode)
 126                 classes = dict( (c._doc_type, c) for c in cls)
 127
 128                 doc = res['doc']
 129                 doc_type = doc['doc_type']
 130                 seq = res['seq']
 131                 c = classes[doc_type]
 132                 yield seq, c.wrap(doc)
 133
 134             since = resp['last_seq']
 135
 136
 137     @staticmethod
 138     def get_cmd_status():
 139         try:
 140             status = CommandStatus.get('episode-slug-status')
 141         except ResourceNotFound:
 142             status = CommandStatus()
 143             status.command = 'Episode-Slug-Update'
 144             status._id = 'episode-slug-status'
 145
 146         return status
 147
 148
 149     @staticmethod
 150     def get_since(status, options):
 151         if options['continue']:
 152             return status.last_seq
 153         else:
 154             return options['since']