Commit f308b2e2 authored by kaiyou's avatar kaiyou

Implement a hook system, resolve items from config and clean the code

parent 5f38718d
import celery
import redis
import json
from amonit import util, config
class State(object):
""" Stores the check state in a redis backend
"""
DEFAULT_STATE = {"up": True, "recurrence": 0}
def __init__(self, redis_url):
self.storage = redis.Redis.from_url(redis_url)
def __getitem__(self, key):
value = self.storage.get(key)
return json.loads(value) if value else State.DEFAULT_STATE.copy()
def __setitem__(self, key, value):
self.storage.set(key, json.dumps(value))
from amonit import util, config, state
class Scheduler(object):
""" Manages celery schedules and handles check responses
"""
def __init__(self, config=None):
def __init__(self, config):
self.config = config
self.state = State(config["general"]["storage"])
self.state = state.State(config["general"]["storage"])
@classmethod
def get_app(cls, conf=None):
......@@ -43,38 +25,27 @@ class Scheduler(object):
for checkid, check in self.config["checks"].items():
app.add_periodic_task(
check["schedule"],
check_run.s(
checkid, check["function"],
check.get("context", {}), check.get("args", {})
)
)
def update(self, checkid, context, result):
""" Handle a status update for a given check
"""
context.update(
recurrence=(context["recurrence"] + 1
if result["up"] == context["up"]
else 0),
**result
)
self.state[checkid] = context
self.notify(checkid, context)
check_run.s(checkid, check["function"],
check.get("context", {}), check.get("args", {})))
def notify(self, checkid, context):
""" Dispatch notifications for a handled status update
"""
for notifierid, notifier in self.config["notifiers"].items():
for criteria in notifier.get("filters", []):
for field, value in criteria.items():
if context.get(field, None) != value:
break
else:
notify_run.s(
notifierid, notifier["function"],
context, notifier.get("args", {})
)()
break
if util.filter_match(context, notifier.get("filters", [])):
args = notifier.get("args", {})
context, args = self.hook("prenotify", context, args)
notify_run.s(notifierid, notifier["function"], context, args)()
def hook(self, trigger, context, io):
""" Handle hooks, io is either the function args or result,
depending if the hook is run pre or post function
"""
for hookid, hook in self.config.get("hooks", {}).items():
if trigger in hook["triggers"]:
function, hook_args = hook["function"], hook.get("args", {})
context, io = util.resolve(function)(context, io, **hook_args)
return context, io
@celery.current_app.task
......@@ -82,12 +53,14 @@ def check_run(checkid, function, context, args):
""" Celery task that runs a single check
"""
print("Running check {}".format(checkid))
stored = celery.current_app.scheduler.state[checkid]
stored.update(context)
stored.update(checkid=checkid, function=function)
result = util.resolve(function)(stored, **args)
celery.current_app.scheduler.update(checkid, stored, result)
return result
scheduler = celery.current_app.scheduler
context, args = scheduler.hook("precheck", context, args)
context = dict(scheduler.state[checkid], **context)
context.update(checkid=checkid, function=function)
result = util.resolve(function)(context, **args)
context, result = scheduler.hook("postcheck", context, result)
scheduler.state[checkid] = context
scheduler.notify(checkid, context)
@celery.current_app.task
......
......@@ -12,17 +12,38 @@ def resolve(function, cache={}):
def render(template, data):
""" Quickly render a jinja template with provided data dict
"""
return jinja2.Template(template).render(**data)
def filter_match(candidate, filters):
""" Apply a list of filters to a candidate dictionary
"""
if type(filters) is list:
return any(
filter_match(candidate, subfilters)
for subfilters in filters
)
elif type(filters) is dict:
return all(
filter_match(candidate.get(key, None), value)
for key, value in filters.items()
)
else:
return candidate == filters
def wrap(function):
""" Wrap a check function to handle simple results and exceptions
"""
def replacement(*args, **kwargs):
try:
result = str(function(*args, **kwargs))
result = function(*args, **kwargs)
if type(result) is not dict or "up" not in result:
result = {
"up": True,
"message": result
"message": str(result)
}
except Exception as error:
result = {
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment