mygpo/maintenance/management/commands/dump-sample.py

   1 from base64 import b64decode
   2 from optparse import make_option
   3 import sys
   4
   5 from couchdb.multipart import write_multipart
   6
   7 from django.core.management.base import BaseCommand
   8
   9 from mygpo.core.models import Podcast
  10 from mygpo.users.models import PodcastUserState, EpisodeUserState, \
  11          Suggestions, User
  12 from mygpo.directory.models import Category
  13 from mygpo.utils import progress
  14 from mygpo.json import json
  15
  16
  17 class Command(BaseCommand):
  18     """
  19     Dumps a Sample of the whole Database that can be used for
  20     testing/development. All objects that are (indirectly) referenced
  21     be the users specified by --user args are dumped.
  22
  23     The dump is similar to a dump of couchdb-python's couchdb-dump and
  24     can be imported by its couchdb-load
  25     """
  26
  27
  28     option_list = BaseCommand.option_list + (
  29         make_option('--user', action='append', type="string", dest='users',
  30             help="User for which related data should be dumped"),
  31     )
  32
  33
  34     def handle(self, *args, **options):
  35
  36         docs = set()
  37
  38         for username in options.get('users', []):
  39             user = User.get_user(username)
  40
  41             # User
  42             docs.add(user._id)
  43
  44             # Suggestions
  45             suggestions = Suggestions.for_user(user)
  46             docs.add(suggestions._id)
  47
  48             # Podcast States
  49             for p_state in PodcastUserState.for_user(user):
  50                 docs.add(p_state._id)
  51
  52                 # Categories
  53                 for tag in p_state.tags:
  54                     c = Category.for_tag(tag)
  55                     if c: docs.add(c._id)
  56
  57                 # Podcast
  58                 podcast = Podcast.get(p_state.podcast)
  59                 docs.add(podcast._id)
  60
  61                 # Categories
  62                 for s in podcast.tags:
  63                     for tag in podcast.tags[s]:
  64                         c = Category.for_tag(tag)
  65                         if c: docs.add(c._id)
  66
  67                 # Episodes
  68                 for episode in podcast.get_episodes():
  69                     docs.add(episode._id)
  70
  71                     # Episode States
  72                     e_state = episode.get_user_state(user)
  73                     if e_state._id:
  74                         docs.add(e_state._id)
  75
  76
  77         db = Podcast.get_db()
  78         docs = sorted(docs)
  79         self.dump(docs, db)
  80
  81
  82     def dump(self, docs, db):
  83
  84         output = sys.stdout
  85         boundary = None
  86         envelope = write_multipart(output, boundary=boundary)
  87         total = len(docs)
  88
  89         for n, docid in enumerate(docs):
  90
  91             doc = db.get(docid, attachments=True)
  92             attachments = doc.pop('_attachments', {})
  93             jsondoc = json.encode(doc)
  94
  95             if attachments:
  96                 parts = envelope.open({
  97                     'Content-ID': doc['_id'],
  98                     'ETag': '"%s"' % doc['_rev']
  99                 })
 100                 parts.add('application/json', jsondoc)
 101
 102                 for name, info in attachments.items():
 103                     content_type = info.get('content_type')
 104                     if content_type is None: # CouchDB < 0.8
 105                         content_type = info.get('content-type')
 106                     parts.add(content_type, b64decode(info['data']), {
 107                         'Content-ID': name
 108                     })
 109                 parts.close()
 110
 111             else:
 112                 envelope.add('application/json', jsondoc, {
 113                     'Content-ID': doc['_id'],
 114                     'ETag': '"%s"' % doc['_rev']
 115                 })
 116
 117             progress(n+1, total, docid, stream=sys.stderr)
 118
 119         envelope.close()