sync/listen.py

import abc
import json
import logging
import os
import urllib.parse

import kombu
from kombu.mixins import ConsumerMixin

from . import log
from . import handlers
from . import tasks

from typing import Any, Dict, Optional

here = os.path.dirname(__file__)

logger = log.get_logger(__name__)


def get_listen_logger(config):
    logger = logging.getLogger(__name__)

    log_dir = os.path.join(config["root"],
                           config["paths"]["logs"])
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)

    basic_formatter = logging.Formatter('[%(asctime)s] %(levelname)s:%(name)s:%(message)s')

    file_handler = logging.handlers.TimedRotatingFileHandler(os.path.join(log_dir, "listen.log"),
                                                             when="D", utc=True)
    file_handler.setFormatter(basic_formatter)

    logger.addHandler(file_handler)
    logger.propagate = False

    return logger


class Listener(ConsumerMixin):
    """Manages a single kombu.Consumer."""
    def __init__(self, conn, exchanges, queues, logger):
        self.connection = conn
        self._callbacks = {item: [] for item in exchanges}
        self._queues = queues
        self.connect_max_retries = 10
        self.logger = logger

    def get_consumers(self, Consumer, channel):
        consumer = Consumer(self._queues, callbacks=[self.on_message], auto_declare=False)
        return [consumer]

    def on_connection_revived(self):
        logger.debug("Connection to %s revived." % self.connection.hostname)

    def add_callback(self, exchange, func):
        if exchange is None:
            raise ValueError("Expected string, got None")

        self._callbacks[exchange].append(func)

    def on_message(self, body, message):
        exchange = message.delivery_info['exchange']
        callbacks = self._callbacks.get(exchange)
        try:
            if callbacks:
                for cb in callbacks:
                    cb(body)
            else:
                raise Exception('received message from unknown exchange: %s' %
                                exchange)
        finally:
            message.ack()


def get_listener(conn, userid, exchanges=None, extra_data=None, logger=None):
    """Obtain a Pulse consumer that can handle received messages.

    Returns a ``Listener`` instance bound to listen to the requested exchanges.
    Callers should use ``add_callback`` to register functions
    that will be called when a message is received.

    The callback functions receive one argument ``body``, the decoded message body.
    """
    queues = []

    if exchanges is None:
        raise ValueError("No exchanges supplied")

    for queue_name, exchange_name, key_name in exchanges:
        queue_name = f'queue/{userid}/{queue_name}'

        exchange = kombu.Exchange(exchange_name, type='topic',
                                  channel=conn)
        exchange.declare(passive=True)

        queue = kombu.Queue(name=queue_name,
                            exchange=exchange,
                            durable=True,
                            routing_key=key_name,
                            exclusive=False,
                            auto_delete=False,
                            channel=conn,
                            extra_data=extra_data)
        queues.append(queue)
        # queue.declare() declares the exchange, which isn't allowed by the
        # server. So call the low-level APIs to only declare the queue itself.
        queue.queue_declare()
        queue.queue_bind()

    return Listener(conn, [item[1] for item in exchanges], queues, logger)


def run_pulse_listener(config: Dict[str, Any]) -> None:
    """
    Configures Pulse connection and triggers events from Pulse messages.

    Connection details are managed at https://pulseguardian.mozilla.org/.
    """
    exchanges = []
    queues = {}
    for queue_name, queue_props in config['pulse'].items():
        if (isinstance(queue_props, dict) and
            set(queue_props.keys()) == {"queue", "exchange", "routing_key"}):
            queues[queue_name] = queue_props

    for queue in queues.values():
        logger.info("Connecting to pulse queue:%(queue)s exchange:%(exchange)s"
                    " route:%(routing_key)s" % queue)
        exchanges.append((queue['queue'],
                          queue['exchange'],
                          queue['routing_key']))

    conn = kombu.Connection(hostname=config['pulse']['host'],
                            port=config['pulse']['port'],
                            ssl=config['pulse']['ssl'],
                            userid=config['pulse']['username'],
                            password=config['pulse']['password'])

    listen_logger = get_listen_logger(config)

    filter_map = {
        'github': GitHubFilter,
        'hgmo': PushFilter,
        'taskcluster-taskgroup': TaskGroupFilter,
        'taskcluster-try-completed': DecisionTaskFilter,
        'taskcluster-try-failed': DecisionTaskFilter,
        'taskcluster-try-exception': DecisionTaskFilter,
        'taskcluster-wptsync-completed': TryTaskFilter,
        'taskcluster-wptsync-failed': TryTaskFilter,
        'taskcluster-wptsync-exception': TryTaskFilter,
    }

    with conn:
        try:
            listener = get_listener(conn,
                                    userid=config['pulse']['username'],
                                    exchanges=exchanges,
                                    logger=listen_logger)
            for queue_name, queue in queues.items():
                queue_filter = filter_map[queue_name](config, listen_logger)
                listener.add_callback(queue['exchange'], queue_filter)

            listener.run()
        except KeyboardInterrupt:
            pass
    return None


class Filter(metaclass=abc.ABCMeta):
    name: Optional[str] = None
    task = tasks.handle

    def __init__(self, config, logger):
        self.config = config
        self.logger = logger

    def __call__(self, body):
        if self.accept(body):
            self.logger.info("Message accepted %s" % self.name)
            self.logger.debug(json.dumps(body))
            self.task.apply_async((self.name, body))

    def accept(self, body: Dict[str, Any]) -> bool:
        raise NotImplementedError


class GitHubFilter(Filter):
    name = "github"
    event_filters = {item: lambda x: True
                     for item in handlers.GitHubHandler.dispatch_event.keys()}
    event_filters["check_run"] = lambda x: x["payload"]["action"] == "completed"
    event_filters["push"] = lambda x: x["payload"]["ref"] == "refs/heads/master"

    def __init__(self, config, logger):
        super().__init__(config, logger)
        repo_path = urllib.parse.urlparse(config["web-platform-tests"]["repo"]["url"]).path
        self.key_filter = "%s/" % repo_path.split("/", 2)[1]

    def accept(self, body):
        return (body["_meta"]["routing_key"].startswith(self.key_filter) and
                body["event"] in self.event_filters and
                self.event_filters[body["event"]](body))


class PushFilter(Filter):
    name = "push"

    def __init__(self, config, logger):
        super().__init__(config, logger)
        self.repos = set(config["gecko"]["repo"].keys())

    def accept(self, body):
        # Check that this has some commits pushed
        if not body["payload"].get("data", {}).get("pushlog_pushes"):
            return False

        repo = body["_meta"]["routing_key"]
        if "/" in repo:
            repo = repo.rsplit("/", 1)[1]
        return repo in self.repos


class TaskGroupFilter(Filter):
    name = "taskgroup"

    def accept(self, body):
        return body.get("taskGroupId") is not None


class DecisionTaskFilter(Filter):
    name = "decision-task"

    def accept(self, body):
        return is_decision_task(body)


class TryTaskFilter(Filter):
    name = "try-task"

    def accept(self, body):
        return not is_decision_task(body)


def is_decision_task(body):
    tags = body.get("task", {}).get("tags", {})
    return (tags.get("kind") == "decision-task" and
            tags.get("createdForUser") == "wptsync@mozilla.com")