From 5a199fec4edd12209355fb51d408b498578a1c69 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C4=93teris=20Caune?= Date: Wed, 30 Dec 2015 21:56:03 +0200 Subject: [PATCH] management commands use self.stdout.write instead of print. Fixes #21 --- hc/api/management/commands/ensuretriggers.py | 6 +- hc/api/management/commands/fillnpings.py | 4 +- hc/api/management/commands/prunechecks.py | 3 +- hc/api/management/commands/prunepings.py | 12 +- hc/api/management/commands/sendalerts.py | 110 +++++++++---------- hc/api/management/commands/sendreports.py | 7 +- hc/api/tests/test_sendalerts.py | 6 +- 7 files changed, 73 insertions(+), 75 deletions(-) diff --git a/hc/api/management/commands/ensuretriggers.py b/hc/api/management/commands/ensuretriggers.py index e5e5a50b..695977fd 100644 --- a/hc/api/management/commands/ensuretriggers.py +++ b/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" diff --git a/hc/api/management/commands/fillnpings.py b/hc/api/management/commands/fillnpings.py index 194aa573..037f292f 100644 --- a/hc/api/management/commands/fillnpings.py +++ b/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!" diff --git a/hc/api/management/commands/prunechecks.py b/hc/api/management/commands/prunechecks.py index 96936ff7..c98f015c 100644 --- a/hc/api/management/commands/prunechecks.py +++ b/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 diff --git a/hc/api/management/commands/prunepings.py b/hc/api/management/commands/prunepings.py index 31edf77f..96a4af5e 100644 --- a/hc/api/management/commands/prunepings.py +++ b/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("---") - print("User: %s" % check.user.email) - print("Check: %s" % check.name) - print("Pruned: %d" % n) + total += n + self.stdout.write("---") + self.stdout.write("User: %s" % check.user.email) + self.stdout.write("Check: %s" % check.name) + self.stdout.write("Pruned: %d" % n) - print("Done.") + return "Done! Pruned %d pings." % total diff --git a/hc/api/management/commands/sendalerts.py b/hc/api/management/commands/sendalerts.py index 080fcfd6..333d1b3e 100644 --- a/hc/api/management/commands/sendalerts.py +++ b/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) diff --git a/hc/api/management/commands/sendreports.py b/hc/api/management/commands/sendreports.py index 0a2ac688..9ef55eb2 100644 --- a/hc/api/management/commands/sendreports.py +++ b/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 diff --git a/hc/api/tests/test_sendalerts.py b/hc/api/tests/test_sendalerts.py index 46462c8f..52c27666 100644 --- a/hc/api/tests/test_sendalerts.py +++ b/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 = []