550 lines
20 KiB
Python
550 lines
20 KiB
Python
"""
|
|
|
|
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
|
|
):
|
|
|
|
"""
|
|
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.
|
|
: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
|
|
|
|
# Done here:
|
|
return config
|
|
|
|
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):
|
|
|
|
"""
|
|
Flushes the buffer entirely.
|
|
:return: None
|
|
"""
|
|
|
|
start_time = time.time()
|
|
self.__producer.flush()
|
|
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
|
|
) -> 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.
|
|
: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
|
|
)
|
|
|
|
# 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
|
|
):
|
|
|
|
"""
|
|
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.
|
|
: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
|
|
|
|
# Done here:
|
|
return config
|
|
|
|
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()
|