Source code for util.rabbitmq

from config import Config
from typing import Literal




class RabbitMQUtil:
    """Utilities work working RabbitMQ."""



    @staticmethod
    def get_connection():
        """Get connection to rabbitmq."""
        from urllib.parse import quote_plus, urlunsplit
        from airflow.models.connection import Connection

        conn = Connection.get_connection_from_secrets(Config.RAD_RMQ_CONN)

        print("Retrieving connection", conn.conn_id)
        scheme = "amqp"
        login = conn.login
        password = conn.password
        netloc = conn.host
        if login is not None and password is not None:
            netloc = f"{quote_plus(login)}:{quote_plus(password)}@{netloc}"
        if conn.port:
            netloc = f"{netloc}:{conn.port}"
        path = f"/{conn.schema}"
        connection_str= urlunsplit((scheme, netloc, path, "", ""))
        print(f"Connecting to {connection_str}")
        return connection_str




    @staticmethod
    def get_channel():
        import pika

        connection = pika.BlockingConnection(
            (pika.URLParameters(RabbitMQUtil.get_connection()))
        )
        channel = connection.channel()
        
        return channel




    @staticmethod
    def publish_messages(exchange: str, queue_name: str, messages: list[dict]):
        import json

        channel = RabbitMQUtil.get_channel()
        for message in messages:
            channel.basic_publish(
                exchange=exchange, routing_key=queue_name, body=json.dumps(message)
            )
            
        print("RabbitMq Sent:", len(messages), "messages")
        channel.close()




    @staticmethod
    def publish_rad_messages(
        queue_name: Literal[
            "address.ingest.regrid", "address.ingest.safegraph", "address.ingest.fabric"
        ],
        messages: list[dict],
    ):
        RabbitMQUtil.publish_messages(Config.RAD_PREFIXES.DISCOVERY_QUEUE, queue_name, messages)