<filename>dnstap_receiver/outputs/output_rabbitmq.py
import asyncio
import logging
import time

from dnstap_receiver.outputs import transform
from dnstap_receiver import statistics

try:
    import pika
    has_pika = True
except:
    has_pika = False

clogger = logging.getLogger("dnstap_receiver.console")


def checking_conf(cfg: dict) -> bool:
    """validate the config"""
    clogger.debug("Output handler: rabbitmq")

    valid_conf = True

    if not has_pika:
        valid_conf = False
        clogger.error("Output handler: rabbitmq: confluent_kafka dependency is missing")

    if cfg["connection"]["username"] is None \
        or cfg["connection"]["password"] is None \
        or cfg["connection"]["host"] is None \
        or cfg["connection"]["port"] is None:
        valid_conf = False
        clogger.error("Output handler: rabbitmq: missing connection details")

    if cfg["queue"]["queue"] is None:
        valid_conf = False
        clogger.error("Output handler: rabbitmq: no queue provided")

    return valid_conf


class RabbitMQ:
    """Class to handle RabbitMQ connections and channel push"""

    def __init__(self, output_cfg: dict) -> None:
        self.cfg = output_cfg
        self.cfg["routing_key"] = output_cfg.get("routing_key", output_cfg["queue"]["queue"])

        self.connection: pika.BlockingConnection = None
        self.channel: pika.adapters.blocking_connection.BlockingChannel = None

        self.credentials = pika.PlainCredentials(
                                self.cfg["connection"]["username"],
                                self.cfg["connection"]["password"]
        )
        self.connection_params = pika.ConnectionParameters(
                                host=self.cfg["connection"]["host"],
                                port=self.cfg["connection"]["port"],
                                credentials=self.credentials
        )
        self.init_connection()

    def init_connection(self) -> None:
        """init connection and channel"""
        if self.connection and self.connection.is_open:
            return

        self.connection = pika.BlockingConnection(self.connection_params)

        self.channel = self.connection.channel()
        self.channel.queue_declare(
                    queue       =   self.cfg["queue"]["queue"],
                    passive     =   self.cfg["queue"]["passive"],
                    durable     =   self.cfg["queue"]["durable"],
                    exclusive   =   self.cfg["queue"]["exclusive"],
                    auto_delete =   self.cfg["queue"]["auto_delete"]
        )


    def publish(self, msg) -> None:
        """publish msg to the channel"""
        for attempt in range(self.cfg['retry-count']):
            try:
                self.init_connection()
                clogger.debug("RabbitMQ publish")
                self.channel.basic_publish(
                                exchange=self.cfg["exchange"],
                                routing_key=self.cfg["routing_key"],
                                body=msg
                )
            except (pika.exceptions.ConnectionClosed,
                    pika.exceptions.StreamLostError,
                    pika.exceptions.ChannelWrongStateError,
                    ConnectionResetError
                    ) as connection_error:
                clogger.debug(connection_error)
                clogger.debug(f"Publish failed, trying to reconnect, attepmt {attempt +1}")
                time.sleep(self.cfg['retry-delay'])
            else:
                break
        else:
            clogger.error(f"Publish failed. Connection error after {self.cfg['retry-count']}")


    def close_connection(self):
        """properly close the connection"""
        if self.connection and self.connection.is_open:
            self.connection.close()



async def handle(output_cfg: dict, queue: asyncio.Queue, _metrics: statistics.Statistics, start_shutdown: asyncio.Event):
    """Connect to rabbit and push the messages from the queue"""

    rabbitmq = RabbitMQ(output_cfg=output_cfg)
    clogger.info("Output handler: rabbitmq: Enabled")
    while not start_shutdown.is_set():
        try:
            tapmsg = await asyncio.wait_for(queue.get(), timeout=0.5)
        except asyncio.TimeoutError:
            continue
        msg = transform.convert_dnstap(fmt=output_cfg["format"], tapmsg=tapmsg)
        rabbitmq.publish(msg)
        queue.task_done()

    # tell producer to shut down
    clogger.info("Output handler: rabbitmq: Triggering producer shutdown")
    rabbitmq.close_connection()

<filename>dnstap_receiver/outputs/output_rabbitmq.py
import asyncio
import logging
import time

from dnstap_receiver.outputs import transform
from dnstap_receiver import statistics

try:
    import pika
    has_pika = True
except:
    has_pika = False

clogger = logging.getLogger("dnstap_receiver.console")


def checking_conf(cfg: dict) -> bool:
    """validate the config"""
    clogger.debug("Output handler: rabbitmq")

    valid_conf = True

    if not has_pika:
        valid_conf = False
        clogger.error("Output handler: rabbitmq: confluent_kafka dependency is missing")

    if cfg["connection"]["username"] is None \
        or cfg["connection"]["password"] is None \
        or cfg["connection"]["host"] is None \
        or cfg["connection"]["port"] is None:
        valid_conf = False
        clogger.error("Output handler: rabbitmq: missing connection details")

    if cfg["queue"]["queue"] is None:
        valid_conf = False
        clogger.error("Output handler: rabbitmq: no queue provided")

    return valid_conf


class RabbitMQ:
    """Class to handle RabbitMQ connections and channel push"""

    def __init__(self, output_cfg: dict) -> None:
        self.cfg = output_cfg
        self.cfg["routing_key"] = output_cfg.get("routing_key", output_cfg["queue"]["queue"])

        self.connection: pika.BlockingConnection = None
        self.channel: pika.adapters.blocking_connection.BlockingChannel = None

        self.credentials = pika.PlainCredentials(
                                self.cfg["connection"]["username"],
                                self.cfg["connection"]["password"]
        )
        self.connection_params = pika.ConnectionParameters(
                                host=self.cfg["connection"]["host"],
                                port=self.cfg["connection"]["port"],
                                credentials=self.credentials
        )
        self.init_connection()

    def init_connection(self) -> None:
        """init connection and channel"""
        if self.connection and self.connection.is_open:
            return

        self.connection = pika.BlockingConnection(self.connection_params)

        self.channel = self.connection.channel()
        self.channel.queue_declare(
                    queue       =   self.cfg["queue"]["queue"],
                    passive     =   self.cfg["queue"]["passive"],
                    durable     =   self.cfg["queue"]["durable"],
                    exclusive   =   self.cfg["queue"]["exclusive"],
                    auto_delete =   self.cfg["queue"]["auto_delete"]
        )


    def publish(self, msg) -> None:
        """publish msg to the channel"""
        for attempt in range(self.cfg['retry-count']):
            try:
                self.init_connection()
                clogger.debug("RabbitMQ publish")
                self.channel.basic_publish(
                                exchange=self.cfg["exchange"],
                                routing_key=self.cfg["routing_key"],
                                body=msg
                )
            except (pika.exceptions.ConnectionClosed,
                    pika.exceptions.StreamLostError,
                    pika.exceptions.ChannelWrongStateError,
                    ConnectionResetError
                    ) as connection_error:
                clogger.debug(connection_error)
                clogger.debug(f"Publish failed, trying to reconnect, attepmt {attempt +1}")
                time.sleep(self.cfg['retry-delay'])
            else:
                break
        else:
            clogger.error(f"Publish failed. Connection error after {self.cfg['retry-count']}")


    def close_connection(self):
        """properly close the connection"""
        if self.connection and self.connection.is_open:
            self.connection.close()



async def handle(output_cfg: dict, queue: asyncio.Queue, _metrics: statistics.Statistics, start_shutdown: asyncio.Event):
    """Connect to rabbit and push the messages from the queue"""

    rabbitmq = RabbitMQ(output_cfg=output_cfg)
    clogger.info("Output handler: rabbitmq: Enabled")
    while not start_shutdown.is_set():
        try:
            tapmsg = await asyncio.wait_for(queue.get(), timeout=0.5)
        except asyncio.TimeoutError:
            continue
        msg = transform.convert_dnstap(fmt=output_cfg["format"], tapmsg=tapmsg)
        rabbitmq.publish(msg)
        queue.task_done()

    # tell producer to shut down
    clogger.info("Output handler: rabbitmq: Triggering producer shutdown")
    rabbitmq.close_connection()

validate the config Class to handle RabbitMQ connections and channel push init connection and channel publish msg to the channel properly close the connection Connect to rabbit and push the messages from the queue # tell producer to shut down