diff --git a/zephyr/management/commands/clear_db.py b/zephyr/management/commands/clear_db.py index e5bc09a05d..fdba5321f9 100644 --- a/zephyr/management/commands/clear_db.py +++ b/zephyr/management/commands/clear_db.py @@ -1,17 +1,9 @@ from django.core.management.base import NoArgsCommand - -from django.contrib.auth.models import User -from zephyr.models import Message, UserProfile, Stream, Recipient, \ - Subscription, Huddle, Realm, UserMessage, Client -from django.contrib.sessions.models import Session +from zephyr.models import clear_database class Command(NoArgsCommand): help = "Clear only tables we change: messages, accounts + sessions" def handle_noargs(self, **options): - for model in [Message, Stream, UserProfile, User, Recipient, - Realm, Subscription, Huddle, UserMessage, Client]: - model.objects.all().delete() - Session.objects.all().delete() - + clear_database() self.stdout.write("Successfully cleared the database.\n") diff --git a/zephyr/management/commands/populate_db.py b/zephyr/management/commands/populate_db.py index 379ca85932..b8dde15f8a 100644 --- a/zephyr/management/commands/populate_db.py +++ b/zephyr/management/commands/populate_db.py @@ -6,14 +6,13 @@ from zephyr.models import Message, UserProfile, Stream, Recipient, Client, \ Subscription, Huddle, get_huddle, Realm, UserMessage, get_user_profile_by_id, \ bulk_create_realms, bulk_create_streams, bulk_create_users, bulk_create_huddles, \ bulk_create_clients, \ - do_send_message, \ + do_send_message, clear_database, \ get_huddle_hash, get_client, do_activate_user from zephyr.lib.parallel import run_parallel from django.db import transaction from django.conf import settings from api.bots import mit_subs_list from zephyr.lib.bulk_create import batch_bulk_create -from django.contrib.sessions.models import Session import simplejson import datetime @@ -99,10 +98,8 @@ class Command(BaseCommand): return if options["delete"]: - for model in [Message, Stream, UserProfile, User, Recipient, - Realm, Subscription, Huddle, UserMessage, Client]: - model.objects.all().delete() - Session.objects.all().delete() + # Start by clearing all the data in our database + clear_database() # Create our two default realms humbug_realm = Realm.objects.create(domain="humbughq.com") diff --git a/zephyr/models.py b/zephyr/models.py index 9e879c895d..38dfbe2c70 100644 --- a/zephyr/models.py +++ b/zephyr/models.py @@ -16,6 +16,7 @@ from zephyr.lib.avatar import gravatar_hash import requests from django.contrib.auth.models import UserManager from django.utils import timezone +from django.contrib.sessions.models import Session @cache_with_key(lambda self: 'display_recipient_dict:%d' % (self.id)) def get_display_recipient(recipient): @@ -616,3 +617,9 @@ def filter_by_subscriptions(messages, user): user_messages.append(message) return user_messages + +def clear_database(): + for model in [Message, Stream, UserProfile, User, Recipient, + Realm, Subscription, Huddle, UserMessage, Client]: + model.objects.all().delete() + Session.objects.all().delete()