management commands use self.stdout.write instead of print. Fixes #21

hc/api/management/commands/ensuretriggers.py

@@ -56,10 +56,10 @@ class Command(BaseCommand):
 
         if connection.vendor == "postgresql":
             _pg(cursor)
-            print("Created PostgreSQL trigger")
+            return "Created PostgreSQL trigger"
         if connection.vendor == "mysql":
             _mysql(cursor)
-            print("Created MySQL trigger")
+            return "Created MySQL trigger"
         if connection.vendor == "sqlite":
             _sqlite(cursor)
-            print("Created SQLite trigger")
+            return "Created SQLite trigger"

hc/api/management/commands/fillnpings.py

@@ -8,6 +8,6 @@ class Command(BaseCommand):
     def handle(self, *args, **options):
         for check in Check.objects.all():
             check.n_pings = Ping.objects.filter(owner=check).count()
-            check.save()
+            check.save(update_fields=("n_pings", ))
 
-        print("Done.")
+        return "Done!"

hc/api/management/commands/prunechecks.py

@@ -10,4 +10,5 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         cutoff = timezone.now() - timedelta(hours=2)
-        Check.objects.filter(user=None, created__lt=cutoff).delete()
+        n, _ = Check.objects.filter(user=None, created__lt=cutoff).delete()
+        return "Done! Pruned %d checks." % n

hc/api/management/commands/prunepings.py

@@ -19,11 +19,13 @@ class Command(BaseCommand):
         checks = checks.annotate(limit=F("user__profile__ping_log_limit"))
         checks = checks.filter(n_pings__gt=F("limit"))
 
+        total = 0
         for check in checks:
             n = check.prune_pings(check.limit)
-            print("---")
+            total += n
-            print("User:   %s" % check.user.email)
+            self.stdout.write("---")
-            print("Check:  %s" % check.name)
+            self.stdout.write("User:   %s" % check.user.email)
-            print("Pruned: %d" % n)
+            self.stdout.write("Check:  %s" % check.name)
+            self.stdout.write("Pruned: %d" % n)
 
-        print("Done.")
+        return "Done! Pruned %d pings." % total

hc/api/management/commands/sendalerts.py

@@ -1,11 +1,9 @@
 import logging
-import sys
 import time
 
 from concurrent.futures import ThreadPoolExecutor
 from django.core.management.base import BaseCommand
 from django.db import connection
-from django.db.models import Q
 from django.utils import timezone
 from hc.api.models import Check
 
@@ -13,72 +11,68 @@ executor = ThreadPoolExecutor(max_workers=10)
 logger = logging.getLogger(__name__)
 
 
-def _stdout(message):
-    sys.stdout.write(message)
-    sys.stdout.flush()
-
-
-def handle_many():
-    """ Send alerts for many checks simultaneously. """
-    query = Check.objects.filter(user__isnull=False)
-
-    now = timezone.now()
-    going_down = Q(alert_after__lt=now, status="up")
-    going_up = Q(alert_after__gt=now, status="down")
-    query = query.filter(going_down | going_up)
-    checks = list(query.iterator())
-    if not checks:
-        return False
-
-    for future in [executor.submit(handle_one, check) for check in checks]:
-        future.result()
-
-    return True
-
-
-def handle_one(check):
-    """ Send an alert for a single check.
-
-    Return True if an appropriate check was selected and processed.
-    Return False if no checks need to be processed.
-
-    """
-    check.status = check.get_status()
-
-    tmpl = "\nSending alert, status=%s, code=%s\n"
-    _stdout(tmpl % (check.status, check.code))
-
-    try:
-        check.send_alert()
-    except:
-        # Catch EVERYTHING. If we crash here, what can happen is:
-        # - the sendalerts command will crash
-        # - supervisor will respawn sendalerts command
-        # - sendalerts will try same thing again, resulting in infinite loop
-        # So instead we catch and log all exceptions, and mark
-        # the checks as paused so they are not retried.
-        logger.error("Could not alert %s" % check.code, exc_info=True)
-        check.status = "paused"
-    finally:
-        check.save()
-        connection.close()
-
-    return True
-
-
 class Command(BaseCommand):
     help = 'Sends UP/DOWN email alerts'
 
+    def handle_many(self):
+        """ Send alerts for many checks simultaneously. """
+        query = Check.objects.filter(user__isnull=False)
+
+        now = timezone.now()
+        going_down = query.filter(alert_after__lt=now, status="up")
+        going_up = query.filter(alert_after__gt=now, status="down")
+        # Don't combine this in one query so Postgres can query using index:
+        checks = list(going_down.iterator()) + list(going_up.iterator())
+        if not checks:
+            return False
+
+        futures = [executor.submit(self.handle_one, check) for check in checks]
+        for future in futures:
+            future.result()
+
+        return True
+
+    def handle_one(self, check):
+        """ Send an alert for a single check.
+
+        Return True if an appropriate check was selected and processed.
+        Return False if no checks need to be processed.
+
+        """
+        check.status = check.get_status()
+
+        tmpl = "\nSending alert, status=%s, code=%s\n"
+        self.stdout.write(tmpl % (check.status, check.code))
+
+        try:
+            check.send_alert()
+        except:
+            # Catch EVERYTHING. If we crash here, what can happen is:
+            # - the sendalerts command will crash
+            # - supervisor will respawn sendalerts command
+            # - sendalerts will try same thing again, resulting in
+            #   infinite loop
+            # So instead we catch and log all exceptions, and mark
+            # the checks as paused so they are not retried.
+            logger.error("Could not alert %s" % check.code, exc_info=True)
+            check.status = "paused"
+        finally:
+            check.save()
+            connection.close()
+
+        return True
+
     def handle(self, *args, **options):
+        self.stdout.write("sendalerts starts up")
 
         ticks = 0
         while True:
-            if handle_many():
+            if self.handle_many():
                 ticks = 0
             else:
                 ticks += 1
 
             time.sleep(1)
-            _stdout(".")
             if ticks % 60 == 0:
-                _stdout("\n")
+                formatted = timezone.now().isoformat()
+                self.stdout.write("-- MARK %s --" % formatted)

hc/api/management/commands/sendreports.py

@@ -16,11 +16,12 @@ def num_pinged_checks(profile):
 
 class Command(BaseCommand):
     help = 'Send due monthly reports'
+    tmpl = "Sending monthly report to %s"
 
     def handle(self, *args, **options):
         # Create any missing profiles
         for u in User.objects.filter(profile__isnull=True):
-            print("Creating profile for %s" % u.email)
+            self.stdout.write("Creating profile for %s" % u.email)
             Profile.objects.for_user(u)
 
         now = timezone.now()
@@ -35,8 +36,8 @@ class Command(BaseCommand):
         sent = 0
         for profile in q:
             if num_pinged_checks(profile) > 0:
-                print("Sending monthly report to %s" % profile.user.email)
+                self.stdout.write(self.tmpl % profile.user.email)
                 profile.send_report()
                 sent += 1
 
-        print("Sent %d reports" % sent)
+        return "Sent %d reports" % sent

hc/api/tests/test_sendalerts.py

@@ -2,14 +2,14 @@ from datetime import datetime
 
 from django.contrib.auth.models import User
 from django.test import TestCase
-from hc.api.management.commands.sendalerts import handle_many
+from hc.api.management.commands.sendalerts import Command
 from hc.api.models import Check
 from mock import patch
 
 
 class SendAlertsTestCase(TestCase):
 
-    @patch("hc.api.management.commands.sendalerts.handle_one")
+    @patch("hc.api.management.commands.sendalerts.Command.handle_one")
     def test_it_handles_few(self, mock):
         alice = User(username="alice")
         alice.save()
@@ -22,7 +22,7 @@ class SendAlertsTestCase(TestCase):
             check.status = "up"
             check.save()
 
-        result = handle_many()
+        result = Command().handle_many()
         assert result, "handle_many should return True"
 
         handled_names = []