Resetting utils subtree.

This commit is contained in:
2025-01-07 12:54:12 +05:30
parent a5c23c05b6
commit a7de497e81
187 changed files with 420 additions and 141093 deletions
-557
View File
@@ -1,557 +0,0 @@
"""
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
@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):
"""
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
@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()