<commit_before>#!/usr/bin/env python3
"""Keepalived notification receiver.

Invoked by keepalived, if the state of VRRP instances changes.
"""
import logging
import os
import sys
import textwrap
import typing as t
from contextlib import closing

import kombu

from hades.agent import create_app
from hades.common.cli import ArgumentParser, common_parser, setup_cli_logging
from hades.common.exc import handles_setup_errors
from hades.config import (
    Config,
    load_config,
    Option,
)
from hades.config.options import (
    HADES_CELERY_NODE_QUEUE,
    HADES_CELERY_NOTIFY_EXCHANGE,
    HADES_CELERY_NOTIFY_EXCHANGE_TYPE,
    HADES_CELERY_ROUTING_KEY_MASTERS_ALL,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE_AUTH,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE_ROOT,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE_UNAUTH,
    HADES_CELERY_RPC_EXCHANGE,
    HADES_CELERY_RPC_EXCHANGE_TYPE,
)

logger = logging.getLogger('hades.bin.vrrp_notify')
INSTANCES: t.Dict[str, t.Type[Option]] = {
    "hades-auth": HADES_CELERY_ROUTING_KEY_MASTERS_SITE_AUTH,
    "hades-root": HADES_CELERY_ROUTING_KEY_MASTERS_SITE_ROOT,
    "hades-unauth": HADES_CELERY_ROUTING_KEY_MASTERS_SITE_UNAUTH,
}


def update_bindings(config: Config, name: str, state: str) -> None:
    app = create_app(config)
    queue_name = config[HADES_CELERY_NODE_QUEUE]
    rpc_exchange = kombu.Exchange(
        config[HADES_CELERY_RPC_EXCHANGE],
        config[HADES_CELERY_RPC_EXCHANGE_TYPE],
    )
    notify_exchange = kombu.Exchange(
        config[HADES_CELERY_NOTIFY_EXCHANGE],
        config[HADES_CELERY_NOTIFY_EXCHANGE_TYPE],
    )
    instance_key = config[INSTANCES[name]]
    bindings = {
        rpc_exchange: {instance_key},
        notify_exchange: {
            config[HADES_CELERY_ROUTING_KEY_MASTERS_ALL],
            config[HADES_CELERY_ROUTING_KEY_MASTERS_SITE],
            instance_key,
        },
    }
    with closing(app.connection(connect_timeout=1)) as connection:
        queue = app.amqp.queues[queue_name]
        bound_queue = queue.bind(connection.default_channel)
        bound_queue.declare()
        if state == 'MASTER':
            for exchange, keys in bindings.items():
                bound_exchange = exchange.bind(connection.default_channel)
                bound_exchange.declare()
                for key in keys:
                    logger.info(
                        "Binding node queue %s to exchange %s "<commit_msg>Try connecting to AMQP only once in VRRP script<commit_after>#!/usr/bin/env python3
"""Keepalived notification receiver.

Invoked by keepalived, if the state of VRRP instances changes.
"""
import logging
import os
import sys
import textwrap
import typing as t

import kombu

from hades.agent import create_app
from hades.common.cli import ArgumentParser, common_parser, setup_cli_logging
from hades.common.exc import handles_setup_errors
from hades.config import (
    Config,
    load_config,
    Option,
)
from hades.config.options import (
    HADES_CELERY_NODE_QUEUE,
    HADES_CELERY_NOTIFY_EXCHANGE,
    HADES_CELERY_NOTIFY_EXCHANGE_TYPE,
    HADES_CELERY_ROUTING_KEY_MASTERS_ALL,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE_AUTH,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE_ROOT,
    HADES_CELERY_ROUTING_KEY_MASTERS_SITE_UNAUTH,
    HADES_CELERY_RPC_EXCHANGE,
    HADES_CELERY_RPC_EXCHANGE_TYPE,
)

logger = logging.getLogger('hades.bin.vrrp_notify')
INSTANCES: t.Dict[str, t.Type[Option]] = {
    "hades-auth": HADES_CELERY_ROUTING_KEY_MASTERS_SITE_AUTH,
    "hades-root": HADES_CELERY_ROUTING_KEY_MASTERS_SITE_ROOT,
    "hades-unauth": HADES_CELERY_ROUTING_KEY_MASTERS_SITE_UNAUTH,
}


def update_bindings(config: Config, name: str, state: str) -> None:
    app = create_app(config)
    queue_name = config[HADES_CELERY_NODE_QUEUE]
    rpc_exchange = kombu.Exchange(
        config[HADES_CELERY_RPC_EXCHANGE],
        config[HADES_CELERY_RPC_EXCHANGE_TYPE],
    )
    notify_exchange = kombu.Exchange(
        config[HADES_CELERY_NOTIFY_EXCHANGE],
        config[HADES_CELERY_NOTIFY_EXCHANGE_TYPE],
    )
    instance_key = config[INSTANCES[name]]
    bindings = {
        rpc_exchange: {instance_key},
        notify_exchange: {
            config[HADES_CELERY_ROUTING_KEY_MASTERS_ALL],
            config[HADES_CELERY_ROUTING_KEY_MASTERS_SITE],
            instance_key,
        },
    }
    with app.connection(connect_timeout=1) as connection:
        connection.ensure_connection(max_retries=0, timeout=1)
        queue = app.amqp.queues[queue_name]
        bound_queue = queue.bind(connection.default_channel)
        bound_queue.declare()
        if state == 'MASTER':
            for exchange, keys in bindings.items():
                bound_exchange = exchange.bind(connection.default_channel)
                bound_exchange.declare()
                for key in keys:
                    logger.info(
                        "Binding node queue %s to exchange %s "
{'fraction_non_alphanumeric': 0.07379185883296571, 'fraction_numerical': 0.00120312813314618, 'mean_word_length': 3.5761467889908256, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6173967', 'n_tokens_mistral': 1829, 'n_tokens_neox': 1797, 'n_words': 282}