""" AUTHOR: Khushal P Soonderji DATE: Friday, 20th Dec., 2024 OBJECTIVE: To work with Kafka in a synchronous way. This is a translation of the wrapper originally made for asynchronous operation. REFERENCES: 01. https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html# DOWNLOADS: N/A """ # ***************************************************************************************************************** # ***** **** # *** IMPORT *** # ***** **** # ***************************************************************************************************************** # To make sibling directories accessible for imports: import sys sys.path.append(".") sys.path.append("..") # To use Kafka: from confluent_kafka import Producer from confluent_kafka import Consumer # For working with JSON strings: from utils_v2.string import json from utils_v2.serialization.json_serializer import JSONSerializer # Data models: from utils_v2.queue.kafka.models.message import ConsumedKafkaMessage # For debugging: from icecream import IceCreamDebugger # To work with datatypes: from typing import List, Literal # To work with date and time: import time # ***************************************************************************************************************** # ***** **** # *** MACROS / ONE-TIME INIT *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** VARIABLES *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** FUNCTIONS *** # ***** **** # ***************************************************************************************************************** # --- Nothing Yet # ***************************************************************************************************************** # ***** **** # *** CLASSES *** # ***** **** # ***************************************************************************************************************** class ProducerKafka: def __init__( self, config: dict, topic: str, serializer = None, debug: bool = True, debug_prefix = "Kafka (P) | " ): """ Create a Kafka Producer. :param config: The configuration as expected by Confluent-Kafka library. :param topic: The topic to produce on. :param serializer: The serializer to use. :param debug: Whether, or not, you want to print the debug strings. :param debug_prefix: The prefix to use while debugging. """ # Initialize the debugger: self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) if not debug: self.__printer.disable() # initialize the Kafka producer: self.__topic = topic self.__config = config self.__producer = None self.__connected = False self.__serializer = serializer or JSONSerializer() def enable_debug(self): self.__printer.enable() def disable_debug(self): self.__printer.disable() @staticmethod def create_config( bootstrap_servers: str | List[str], security_protocol: Literal["PLAINTEXT", "SSL"] = "PLAINTEXT", ca_file: str | None = None, cert_file: str | None = None, key_file: str | None = None, client_id: str | int | None = None, acks: int = 0, retries: int = 1, linger_ms: int = 0, misc_json: dict = None ): """ Creates the config required for Confluent-Kafka's library. :param bootstrap_servers: The addresses of the Kafka brokers. :param security_protocol: What sort of security protocol to use. :param ca_file: Needed for 'SSL' security protocol. :param cert_file: Needed for 'SSL' security protocol. :param key_file: Needed for 'SSL' security protocol. :param client_id: An identifier for one producer. Useful for debugging later. :param acks: The no. of acknowledgements needed by the producer to consider the message as delivered. '0' means fire-and-forget, 'all' or '-1' means that every in-sync replica must acknowledge. It could be any no. :param retries: The no. of times the producer mut try to send the message in case of a failure. :param linger_ms: The time in milliseconds to wait before sending a batch of messages to the broker. :param misc_json: Any other custom configuration that is specific to the Confluent Kafka library. :return: The dictionary that needs to be passed as the 'conf' param when creating the producer. """ # Start with the bare minimum: if not isinstance(bootstrap_servers, list): bootstrap_servers = [bootstrap_servers] config = { "bootstrap.servers": ",".join(bootstrap_servers), "security.protocol": security_protocol, "acks": acks, "retries": retries, "linger.ms": linger_ms } # Add the SSL security details: if security_protocol == "SSL": config["ssl.ca.location"] = ca_file config["ssl.certificate.location"] = cert_file config["ssl.key.location"] = key_file # Add an identifier for debugging: if client_id: config["client.id"] = client_id # Add the elements of the misc. JSON: if misc_json: for k, v in misc_json.items(): config[k] = v # Done here: return config @property def client(self) -> Producer: return self.__producer def connect(self) -> bool: """ Connects to the Kafka server if not connected. :return: True or False based on the success of the operation. """ if not self.__connected: try: self.__producer = Producer(self.__config) self.__connected = True except Exception as exception: self.__printer(exception) return self.__connected def ensure_connection(self): """ Connects to the Kafka server if not connected. :return: True or False based on the success of the operation. """ if not self.__connected: self.connect() return self.__connected def flush(self, timeout: float = None): """ Flushes the buffer entirely. :param timeout: The timeout value in seconds. :return: None """ start_time = time.time() self.__producer.flush(timeout = timeout) message = f"Producer flushed in {time.time() - start_time:.5f} second(s)." self.__printer(message) def close(self): """ Terminates the connection. :return: None. """ if self.__connected: try: self.flush() self.__printer("Producer closed!") self.__connected = False except Exception as exception: self.__printer(exception) return not self.__connected def produce( self, value, key = None, topic = None, encoding = "utf-8", callback = None, poll_immediately: bool = False ) -> bool: """ Sends one message to the Kafka server on the topic that has been set for this instance. :param value: The message to send. :param key: The key to use when you want the messages to follow an order. :param topic: A custom topic for this message, else the topic defined during the creation of this instance will be used by default. :param encoding: The encoding format. :param callback: The function to call for delivery reports. :param poll_immediately: Whether, or not, you would like to run 'poll' method immediately after the produce method. :return: True or False based on the success of the operation. """ # Ensure connectivity to the server. # If not connected, return with failure immediately. if not self.ensure_connection(): return False try: # Send the message: self.__producer.produce( topic = topic or self.__topic, value = self.__serializer.serialize(data = value, encoding = encoding), key = key, callback = callback ) if poll_immediately: self.__producer.poll(0) # Return with success if no exception occurred: return True # Return with failure if something went wrong: except Exception as exception: self.__printer(exception, self.__topic, type(value), value) return False # --------------------------------------------------------------------------------------------------------------------- class ConsumerKafka: def __init__( self, config: dict, topic: str, serializer = None, debug = True, debug_prefix = "Kafka (C) | ", **kwargs ): """ Create a Kafka Consumer. :param topic: the topic to consume on. :param serializer: The serializer to use. :param debug: Whether, or not, you want to print the debug strings. :param debug_prefix: The prefix to use while debugging. :param kwargs: Any configuration parameters for the Kafka instances. """ # Initialize the debugger: self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) if not debug: self.__printer.disable() # initialize the Kafka consumer: self.__config = config self.__topic = topic self.__kwargs = kwargs self.__consumer = None self.__connected = False self.__serializer = serializer or JSONSerializer() def enable_debug(self): self.__printer.enable() def disable_debug(self): self.__printer.disable() @staticmethod def create_config( bootstrap_servers: str | List[str], group_id: str = "default", auto_offset_reset: Literal["latest", "earliest"] = "latest", security_protocol: Literal["PLAINTEXT", "SSL"] = "PLAINTEXT", ca_file: str | None = None, cert_file: str | None = None, key_file: str | None = None, client_id: str | int | None = None, misc_json: dict = None ): """ Creates the config required for Confluent-Kafka's library. :param bootstrap_servers: The addresses of the Kafka brokers. :param group_id: When a set of consumers are working on one topic in a group such that you want only one of them to read a particular message. :param auto_offset_reset: Use this to influence the behaviour of how the Kafka consumer will read messages when it first connects to the broker. It could either want to read the earliest (oldest) messages or the latest (newest) messages from the queue. :param security_protocol: What sort of security protocol to use. :param ca_file: Needed for 'SSL' security protocol. :param cert_file: Needed for 'SSL' security protocol. :param key_file: Needed for 'SSL' security protocol. :param client_id: An identifier for one producer. Useful for debugging later. :param misc_json: Any other custom configuration that is specific to the Confluent Kafka library. :return: The dictionary that needs to be passed as the 'conf' param when creating the producer. """ # Start with the bare minimum: if not isinstance(bootstrap_servers, list): bootstrap_servers = [bootstrap_servers] config = { "bootstrap.servers": ",".join(bootstrap_servers), "security.protocol": security_protocol } # Add the SSL security details: if security_protocol == "SSL": config["ssl.ca.location"] = ca_file config["ssl.certificate.location"] = cert_file config["ssl.key.location"] = key_file # Add an identifier for debugging: if client_id: config["client.id"] = client_id # Consumer-specific: config["group.id"] = group_id config["auto.offset.reset"] = auto_offset_reset # Add the elements of the misc. JSON: if misc_json: for k, v in misc_json.items(): config[k] = v # Done here: return config @property def client(self) -> Consumer: return self.__consumer def connect(self): """ Connects to the Kafka server if not connected. :return: True or False based on the success of the operation. """ if not self.__connected: try: if not isinstance(self.__topic, list): self.__topic = [self.__topic] self.__consumer = Consumer(self.__config) self.__consumer.subscribe(self.__topic) self.__connected = True except Exception as exception: self.__printer(exception) return self.__connected def ensure_connection(self): """ Connects to the Kafka server if not connected. :return: True or False based on the success of the operation. """ if not self.__connected: self.connect() return self.__connected def close(self): """ Terminates the connection. :return: None. """ if self.__connected: try: self.__consumer.close() self.__printer("Consumer closed!") self.__connected = False except Exception as exception: self.__printer(exception) def __consume_one( self, timeout = 0.05, encoding = "utf-8" ) -> dict | None: """ Here we consume exactly one message from the broker. If we need multiple messages (a batch of messages), we call this method as many times as needed. :param timeout: The time in seconds to wait for retrieval. :param encoding: The encoding to use. :return: A dict that holds the details of the message. Null if no message was fetched. """ message = self.__consumer.poll(timeout = timeout) if message is None or message.error(): return None else: return ConsumedKafkaMessage.from_confluent_kafka( message = message, deserializer = lambda x: self.__serializer.deserialize(x, encoding = encoding) ) # message_value = message.value() # if message_value is not None: message_value = self.__serializer.deserialize(message_value, encoding = encoding) # return { # "topic": message.topic(), # "partition": message.partition(), # "offset": message.offset(), # "key": message.key().decode("utf-8") if message.key() else None, # "value": message_value, # "timestamp": message.timestamp()[1], # "headers": message.headers() # } def consume(self, count = 1, timeout = 0.05, encoding = "utf-8"): """ Get messages from the Kafka broker. :param count: The number of messages to get from the Kafka server. :param timeout: The time in seconds to wait for retrieval. :param encoding: The encoding to use. :return: The messages that were received. If no messages are available, an empty list will be returned. """ # Make some variables: start_time = time.time() messages = [] # Ensure connectivity to the server. # If not connected, return with failure immediately. if not self.ensure_connection(): return messages # Get into an indefinite loop. while True: # If timed-out, break out of the loop: elapsed_time = time.time() - start_time if elapsed_time > timeout: break # Get the next message: message = self.__consume_one( timeout = timeout - elapsed_time, encoding = encoding ) # If a message was fetched in the timeout, append it to the list of messages. # Break out of the loop if you have reached the needed no. of messages: if message: messages.append(message) if len(messages) >= count: break # Done here: return messages # ***************************************************************************************************************** # ***** **** # *** MAIN PROGRAM *** # ***** **** # ***************************************************************************************************************** if __name__ == "__main__": import datetime import time def producer_test( count: int = 10, interval: float = 1.0 ): # Create and connect the producer: producer = ProducerKafka( topic = "kft_file_upload", config = ProducerKafka.create_config( bootstrap_servers = "del.ditscentre.in:9092", security_protocol = "SSL", ca_file = r"../../creds/kafka/cert_authority.pem", cert_file = r"../../creds/kafka/fullchain.pem", key_file = r"../../creds/kafka/privkey.pem", client_id = 123 ) ) producer.connect() print("PRODUCER READY!") # Send the message a number of times: for message_no in range(count): my_message = { "ts": datetime.datetime.now().timestamp(), "msgNo": message_no, "payload": { "name": "Bhopli", "color": "orange", "age": "just a baby", } } success = producer.produce(my_message) print("PRODUCED:", success) time.sleep(interval) # Ensure a graceful close: producer.close() def consumer_test(): # Create and connect the producer: consumer = ConsumerKafka( topic = "tick-listners", config = ConsumerKafka.create_config( bootstrap_servers = "del.ditscentre.in:9092", group_id = "test-group", security_protocol = "SSL", ca_file = r"../../../../creds/kafka/cert_authority.pem", cert_file = r"../../../../creds/kafka/fullchain.pem", key_file = r"../../../../creds/kafka/privkey.pem", ) ) consumer.connect() print("CONSUMER READY!") for _ in range(500): messages = consumer.consume(count = 3, timeout = 2.5) print(json.to_string(messages, default = str)) consumer.close() # Run the test code: # producer_test(count = 5) consumer_test()