Cut out duplicated code between clear_database and populate_db.

(imported from commit 4bb44c20d2f0477d8a3db72788f82b0cf8572b2d)

zephyr/management/commands/clear_db.py

@@ -1,17 +1,9 @@
 from django.core.management.base import NoArgsCommand
-
+from zephyr.models import clear_database
-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
 
 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,
+        clear_database()
-                      Realm, Subscription, Huddle, UserMessage, Client]:
-            model.objects.all().delete()
-        Session.objects.all().delete()
-
         self.stdout.write("Successfully cleared the database.\n")

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,
+            # Start by clearing all the data in our database
-                          Realm, Subscription, Huddle, UserMessage, Client]:
+            clear_database()
-                model.objects.all().delete()
-            Session.objects.all().delete()
 
             # Create our two default realms
             humbug_realm = Realm.objects.create(domain="humbughq.com")

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()