From cd383bd4c4b51c31d7fc931f880eeae9c13179ef Mon Sep 17 00:00:00 2001 From: "Marcel R." Date: Fri, 3 Jan 2025 16:24:46 +0100 Subject: [PATCH] Add mattermost contrib. --- docs/contrib/index.rst | 1 + docs/contrib/mattermost.rst | 19 +++++ law.cfg.example | 49 ++++++++++++ law/cli/completion.sh | 2 +- law/contrib/mattermost/__init__.py | 12 +++ law/contrib/mattermost/config.py | 25 ++++++ law/contrib/mattermost/notification.py | 104 +++++++++++++++++++++++++ law/contrib/mattermost/parameter.py | 70 +++++++++++++++++ law/contrib/slack/notification.py | 1 + 9 files changed, 282 insertions(+), 1 deletion(-) create mode 100644 docs/contrib/mattermost.rst create mode 100644 law/contrib/mattermost/__init__.py create mode 100644 law/contrib/mattermost/config.py create mode 100644 law/contrib/mattermost/notification.py create mode 100644 law/contrib/mattermost/parameter.py diff --git a/docs/contrib/index.rst b/docs/contrib/index.rst index c6507d11..f3e800ba 100644 --- a/docs/contrib/index.rst +++ b/docs/contrib/index.rst @@ -41,6 +41,7 @@ The following example shows how a package (e.g. the :py:mod:`~law.docker` packag keras lsf matplotlib + mattermost mercurial numpy pandas diff --git a/docs/contrib/mattermost.rst b/docs/contrib/mattermost.rst new file mode 100644 index 00000000..eb52354c --- /dev/null +++ b/docs/contrib/mattermost.rst @@ -0,0 +1,19 @@ +mattermost +========== + +.. automodule:: law.mattermost + +.. contents:: + + +Function ``notify_mattermost`` +------------------------------ + +.. autofunction:: notify_mattermost + + +Class ``NotifyMattermostParameter`` +----------------------------------- + +.. autoclass:: NotifyMattermostParameter + :members: diff --git a/law.cfg.example b/law.cfg.example index 7759651a..403f2f06 100644 --- a/law.cfg.example +++ b/law.cfg.example @@ -1029,6 +1029,55 @@ ; Type: ``str``, ``None`` ; Default: ``None`` +; mattermost_hook_url +; Desciption: Address of the mattermost webhook to use. For more info, see +; https://developers.mattermost.com/integrate/webhooks/incoming. +; Type: ``str``, ``None`` +; Default: ``None`` + +; mattermost_header +; Desciption: A custom header to add to mattermost notifications. Not used when empty. +; Type: ``str``, ``None`` +; Default: ``None`` + +; mattermost_channel +; Desciption: A custom channel name to send the notification to. When empty, the default configured +; for the webhook is used. +; Type: ``str``, ``None`` +; Default: ``None`` + +; mattermost_user +; Desciption: A custom user name to send the notification with. When empty, the default configured +; for the webhook is used. +; Type: ``str``, ``None`` +; Default: ``None`` + +; mattermost_mention_user +; Desciption: If set, the user with that name is mentioned. +; Type: ``str``, ``None`` +; Default: ``None`` + +; mattermost_icon_url +; Desciption: The address of a custom icon to use. When empty, the default configured for the +; webhook is used. +; Type: ``str``, ``None`` +; Default: "https://media.githubusercontent.com/media/riga/law/refs/heads/master/assets/logo_profile.png" + +; mattermost_icon_emoji +; Desciption: Name of a custom emoji to use instead of the icon url. +; Type: ``str``, ``None`` +; Default: ``None`` + +; mattermost_success_emoji +; Desciption: Emoji to append to the status text in case of success. +; Type: ``str``, ``None`` +; Default: ":tada:" + +; mattermost_failure_emoji +; Desciption: Emoji to append to the status text in case of a failure. +; Type: ``str``, ``None`` +; Default: ":rotating_light:" + ; --- bash_sandbox section ------------------------------------------------------------------------- diff --git a/law/cli/completion.sh b/law/cli/completion.sh index f9217969..dc5910d3 100755 --- a/law/cli/completion.sh +++ b/law/cli/completion.sh @@ -190,7 +190,7 @@ _law_complete() { # complete the "location" subcommand elif [ "${sub_cmd}" = "location" ]; then local words="help" - local contribs="arc awkward cms coffea docker dropbox gfal git glite hdf5 htcondor ipython keras lsf matplotlib mercurial numpy pandas profiling pyarrow rich root singularity slack slurm tasks telegram tensorflow wlcg" + local contribs="arc awkward cms coffea docker dropbox gfal git glite hdf5 htcondor ipython keras lsf matplotlib mattermost mercurial numpy pandas profiling pyarrow rich root singularity slack slurm tasks telegram tensorflow wlcg" local inp="${cur##-}" inp="${inp##-}" COMPREPLY=( $( compgen -W "$( echo ${words} )" -P "--" -- "${inp}" ) $( compgen -W "$( echo ${contribs} )" -- "${inp}" ) ) diff --git a/law/contrib/mattermost/__init__.py b/law/contrib/mattermost/__init__.py new file mode 100644 index 00000000..3fbfb212 --- /dev/null +++ b/law/contrib/mattermost/__init__.py @@ -0,0 +1,12 @@ +# coding: utf-8 +# flake8: noqa + +""" +Mattermost contrib functionality. +""" + +__all__ = ["notify_mattermost", "NotifyMattermostParameter"] + +# provisioning imports +from law.contrib.mattermost.notification import notify_mattermost +from law.contrib.mattermost.parameter import NotifyMattermostParameter diff --git a/law/contrib/mattermost/config.py b/law/contrib/mattermost/config.py new file mode 100644 index 00000000..8d8b2b23 --- /dev/null +++ b/law/contrib/mattermost/config.py @@ -0,0 +1,25 @@ +# coding: utf-8 + +""" +Function returning the config defaults of the mattermost package. +""" + +from __future__ import annotations + +from law._types import Any + + +def config_defaults(default_config: dict) -> dict[str, dict[str, Any]]: + return { + "notifications": { + "mattermost_hook_url": None, + "mattermost_header": None, + "mattermost_channel": None, + "mattermost_user": None, + "mattermost_mention_user": None, + "mattermost_icon_url": "https://media.githubusercontent.com/media/riga/law/refs/heads/master/assets/logo_profile.png", # noqa + "mattermost_icon_emoji": None, + "mattermost_success_emoji": ":tada:", + "mattermost_failure_emoji": ":rotating_light:", + }, + } diff --git a/law/contrib/mattermost/notification.py b/law/contrib/mattermost/notification.py new file mode 100644 index 00000000..48acf55d --- /dev/null +++ b/law/contrib/mattermost/notification.py @@ -0,0 +1,104 @@ +# coding: utf-8 + +""" +Mattermost notifications. +""" + +from __future__ import annotations + +__all__ = ["notify_mattermost"] + +import threading +import traceback + +from law.config import Config +from law.util import escape_markdown +from law.logger import get_logger +from law._types import Any + + +logger = get_logger(__name__) + + +def notify_mattermost( + title: str, + content: str | dict[str, Any], + hook_url: str | None = None, + channel: str | None = None, + user: str | None = None, + mention_user: str | None = None, + icon_url: str | None = None, + icon_emoji: str | None = None, + **kwargs, +) -> bool: + """ + Sends a mattermost notification and returns *True* on success. The communication with the + mattermost API might have some delays and is therefore handled by a thread. The format of the + notification depends on *content*. If it is a string, a simple text notification is sent. + Otherwise, it should be a dictionary whose fields are formatted as key-value pairs. + """ + cfg = Config.instance() + + # get default settings + if not hook_url: + hook_url = cfg.get_expanded("notifications", "mattermost_hook_url") + if not channel: + channel = cfg.get_expanded("notifications", "mattermost_channel") + if not user: + user = cfg.get_expanded("notifications", "mattermost_user") + if not mention_user: + mention_user = cfg.get_expanded("notifications", "mattermost_mention_user") + if not icon_url: + icon_url = cfg.get_expanded("notifications", "mattermost_icon_url") + if not icon_emoji: + icon_emoji = cfg.get_expanded("notifications", "mattermost_icon_emoji") + + if not hook_url: + logger.warning(f"cannot send Mattermost notification, hook_url ({hook_url}) empty") + return False + + # append the user to mention to the title + # unless explicitly set to empty string + mention_text = "" + if mention_user: + mention_text = " (@{})".format(escape_markdown(mention_user.lstrip("@"))) + + # request data for the API call + request_data = {} + if channel: + request_data["channel"] = channel + if user: + request_data["username"] = user + if icon_url: + request_data["icon_url"] = icon_url + if icon_emoji: + request_data["icon_emoji"] = icon_emoji + + # standard or attachment content? + request_data["text"] = f"{title}{mention_text}\n\n" + if isinstance(content, str): + request_data["text"] += content + else: + for k, v in content.items(): + request_data["text"] += f"{k}: {v}\n" + + # extend by arbitrary kwargs + request_data.update(kwargs) + + # threaded, non-blocking API communication + thread = threading.Thread(target=_notify_mattermost, args=(hook_url, request_data)) + thread.start() + + return True + + +def _notify_mattermost(hook_url: str, request_data: dict[str, Any]) -> None: + import requests # type: ignore[import-untyped] + + try: + res = requests.post(hook_url, json=request_data) + if not res.ok: + logger.warning(f"unsuccessful Mattermost API call: {res}") + except Exception as e: + t = traceback.format_exc() + logger.warning(f"could not send Mattermost notification: {e}\n{t}") diff --git a/law/contrib/mattermost/parameter.py b/law/contrib/mattermost/parameter.py new file mode 100644 index 00000000..428893a5 --- /dev/null +++ b/law/contrib/mattermost/parameter.py @@ -0,0 +1,70 @@ +# coding: utf-8 + +""" +Mattermost related parameters. +""" + +from __future__ import annotations + +__all__ = ["NotifyMattermostParameter"] + +from collections import OrderedDict + +from law.config import Config +from law.parameter import NotifyParameter +from law.contrib.mattermost.notification import notify_mattermost +from law._types import Any + + +class NotifyMattermostParameter(NotifyParameter): + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + self.description: str + if not self.description: + self.description = ( + "when true, and the task's run method is decorated with law.decorator.notify, " + "a Mattermost notification is sent once the task finishes" + ) + + @classmethod + def notify(cls, success: bool, title: str, content: dict[str, Any], **kwargs) -> bool: + content = OrderedDict(content) + + # overwrite title + cfg = Config.instance() + header = cfg.get_expanded("notifications", "mattermost_header") + task_block = f"```\n{content['Task']}\n```" + title = f"{header}\n{task_block}" if header else task_block + del content["Task"] + + # markup for traceback + if "Traceback" in content: + content["Traceback"] = f"\n```\n{content['Traceback']}\n```" + + # prepend the status text to the message content + cfg = Config.instance() + status_text = "success" if success else "failure" + status_emoji = cfg.get_expanded("notifications", f"mattermost_{status_text}_emoji") + if status_emoji: + status_text += " " + status_emoji + content["Status"] = status_text + content.move_to_end("Status", last=False) + + # highlight last message + if "Last message" in content: + content["Last message"] = f"`{content['Last Message']}`" + + # highlight keys + content = content.__class__((f"**{k}**", v) for k, v in content.items()) + + # send the notification + return notify_mattermost(title, content, **kwargs) + + def get_transport(self) -> dict[str, Any]: + return { + "func": self.notify, + "raw": True, + "colored": False, + } diff --git a/law/contrib/slack/notification.py b/law/contrib/slack/notification.py index a6c98c2d..f224a668 100644 --- a/law/contrib/slack/notification.py +++ b/law/contrib/slack/notification.py @@ -19,6 +19,7 @@ from law.logger import get_logger from law._types import Any, ModuleType + logger = get_logger(__name__)