zulip/zerver/management/commands/dump_pointers.py

from __future__ import absolute_import

from optparse import make_option
from django.core.management.base import BaseCommand
from zerver.models import UserProfile, Message, UserMessage, \
    get_user_profile_by_email
from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
import ujson

def dump():
    pointers = []
    for u in UserProfile.objects.all():
        pointer = u.pointer
        if pointer != -1:
            pub_date = Message.objects.get(id=pointer).pub_date
            pointers.append((u.email, datetime_to_timestamp(pub_date)))
        else:
            pointers.append((u.email, -1))
    file("dumped-pointers", "w").write(ujson.dumps(pointers) + "\n")

def restore(change):
    for (email, timestamp) in ujson.loads(file("dumped-pointers").read()):
        try:
            u = get_user_profile_by_email(email)
        except UserProfile.DoesNotExist:
            print "Skipping...", email
            continue
        if timestamp == -1:
            pointer = -1
        else:
            try:
                pointer = UserMessage.objects.filter(user_profile=u,
                    message__pub_date__gte=timestamp_to_datetime(timestamp)).order_by("message")[0].message_id
            except IndexError:
                print "Alert...", email, timestamp
                continue
        if change:
            u.pointer = pointer
            u.save(update_fields=["pointer"])

class Command(BaseCommand):
    option_list = BaseCommand.option_list + (
        make_option('--restore', default=False, action='store_true'),
        make_option('--dry-run', '-n', default=False, action='store_true'),)

    def handle(self, *args, **options):
        if options["restore"]:
            restore(change=not options['dry_run'])
        else:
            dump()