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

This commit is contained in:
Pēteris Caune 2015-12-30 21:56:03 +02:00
parent b7fcaac8ae
commit 5a199fec4e
7 changed files with 73 additions and 75 deletions

View File

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

View File

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

View File

@ -10,4 +10,5 @@ class Command(BaseCommand):
def handle(self, *args, **options): def handle(self, *args, **options):
cutoff = timezone.now() - timedelta(hours=2) 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

View File

@ -19,11 +19,13 @@ class Command(BaseCommand):
checks = checks.annotate(limit=F("user__profile__ping_log_limit")) checks = checks.annotate(limit=F("user__profile__ping_log_limit"))
checks = checks.filter(n_pings__gt=F("limit")) checks = checks.filter(n_pings__gt=F("limit"))
total = 0
for check in checks: for check in checks:
n = check.prune_pings(check.limit) 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

View File

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

View File

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

View File

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