Squashed 'utils_v2/' content from commit f03179d
git-subtree-dir: utils_v2 git-subtree-split: f03179d339e69fdcdff2f0cc06e1f352024c55d3
This commit is contained in:
yatmesh
@@ -0,0 +1,523 @@
|
||||
"""
|
||||
|
||||
AUTHOR:
|
||||
|
||||
Khushal P Soonderji
|
||||
|
||||
DATE:
|
||||
|
||||
Saturday, 13th Jul, 2024
|
||||
|
||||
OBJECTIVE:
|
||||
|
||||
To be able to work with keys
|
||||
|
||||
REFERENCES:
|
||||
|
||||
N/A
|
||||
|
||||
DOWNLOADS:
|
||||
|
||||
N/A
|
||||
|
||||
"""
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** IMPORT ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# To make sibling directories accessible for imports:
|
||||
import sys
|
||||
sys.path.append(".")
|
||||
sys.path.append("..")
|
||||
|
||||
# To use Kafka:
|
||||
from aiokafka import AIOKafkaProducer
|
||||
from aiokafka import AIOKafkaConsumer
|
||||
|
||||
# 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
|
||||
|
||||
# For SSL security:
|
||||
import ssl
|
||||
|
||||
# For asynchronous activities:
|
||||
import asyncio
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** MACROS / ONE-TIME INIT ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# --- Nothing Yet
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** VARIABLES ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# --- Nothing Yet
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** FUNCTIONS ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
def get_ssl_context(
|
||||
ca_file,
|
||||
cert_file,
|
||||
key_file
|
||||
):
|
||||
|
||||
"""
|
||||
Generate the SSL context to use with the Kafka instances.
|
||||
:param ca_file: The Certificate Authority file as a path to a local file.
|
||||
:param cert_file: The Certificate file as a path to a local file.
|
||||
:param key_file: The Key file as a path to a local file.
|
||||
:return: The SSL context instance as a path to a local file.
|
||||
"""
|
||||
|
||||
ssl_context = ssl.create_default_context()
|
||||
ssl_context.load_verify_locations(ca_file)
|
||||
ssl_context.load_cert_chain(certfile = cert_file, keyfile = key_file)
|
||||
return ssl_context
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** CLASSES ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
class ProducerKafka:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
topic,
|
||||
serializer = None,
|
||||
debug = True,
|
||||
debug_prefix = "Kafka (P) | ",
|
||||
**kwargs
|
||||
):
|
||||
|
||||
"""
|
||||
Create a Kafka Producer.
|
||||
: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.
|
||||
: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 producer:
|
||||
self.__topic = topic
|
||||
self.__kwargs = kwargs
|
||||
self.__producer = None
|
||||
self.__connected = False
|
||||
self.__serializer = serializer or JSONSerializer()
|
||||
|
||||
# For establishing connection:
|
||||
self.__exclusive_semaphore = asyncio.Semaphore(1)
|
||||
|
||||
def enable_debug(self):
|
||||
self.__printer.enable()
|
||||
|
||||
def disable_debug(self):
|
||||
self.__printer.disable()
|
||||
|
||||
async def connect(self):
|
||||
|
||||
"""
|
||||
Connects to the Kafka server if not connected.
|
||||
:return: True or False based on the success of the operation.
|
||||
"""
|
||||
|
||||
async with self.__exclusive_semaphore:
|
||||
if not self.__connected:
|
||||
try:
|
||||
self.__producer = AIOKafkaProducer(**self.__kwargs)
|
||||
await self.__producer.start()
|
||||
self.__connected = True
|
||||
except Exception as exception: self.__printer(exception)
|
||||
return self.__connected
|
||||
|
||||
async 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: await self.connect()
|
||||
return self.__connected
|
||||
|
||||
async def close(self):
|
||||
|
||||
"""
|
||||
Terminates the connection.
|
||||
:return: None.
|
||||
"""
|
||||
|
||||
if self.__connected:
|
||||
try:
|
||||
await self.__producer.stop()
|
||||
self.__printer("Producer closed!")
|
||||
self.__connected = False
|
||||
except Exception as exception: self.__printer(exception)
|
||||
|
||||
async def produce(self, value, key = None, topic = None, encoding = "utf-8"):
|
||||
|
||||
"""
|
||||
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.
|
||||
: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 await self.ensure_connection(): return False
|
||||
|
||||
try:
|
||||
|
||||
# Send the message:
|
||||
await self.__producer.send_and_wait(
|
||||
topic = topic or self.__topic,
|
||||
value = self.__serializer.serialize(data = value, encoding = encoding),
|
||||
key = key
|
||||
)
|
||||
|
||||
# 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,
|
||||
topic,
|
||||
serializer = None,
|
||||
debug = True,
|
||||
debug_prefix = "Kafka (C) | ",
|
||||
**kwargs
|
||||
):
|
||||
|
||||
"""
|
||||
Create a Kafka Consumer.
|
||||
:param topic: The topic to consumer 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 producer:
|
||||
self.__topic = topic
|
||||
self.__kwargs = kwargs
|
||||
self.__consumer = None
|
||||
self.__connected = False
|
||||
self.__serializer = serializer or JSONSerializer()
|
||||
|
||||
# For establishing connection:
|
||||
self.__exclusive_semaphore = asyncio.Semaphore(1)
|
||||
|
||||
def enable_debug(self):
|
||||
self.__printer.enable()
|
||||
|
||||
def disable_debug(self):
|
||||
self.__printer.disable()
|
||||
|
||||
async def connect(self):
|
||||
|
||||
"""
|
||||
Connects to the Kafka server if not connected.
|
||||
:return: True or False based on the success of the operation.
|
||||
"""
|
||||
|
||||
async with self.__exclusive_semaphore:
|
||||
if not self.__connected:
|
||||
try:
|
||||
self.__consumer = AIOKafkaConsumer(self.__topic, **self.__kwargs)
|
||||
await self.__consumer.start()
|
||||
self.__connected = True
|
||||
except Exception as exception: self.__printer(exception)
|
||||
return self.__connected
|
||||
|
||||
async 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: await self.connect()
|
||||
return self.__connected
|
||||
|
||||
async def close(self):
|
||||
|
||||
"""
|
||||
Terminates the connection.
|
||||
:return: None.
|
||||
"""
|
||||
|
||||
if self.__connected:
|
||||
try:
|
||||
await self.__consumer.stop()
|
||||
self.__printer("Consumer closed!")
|
||||
self.__connected = False
|
||||
except Exception as exception: self.__printer(exception)
|
||||
|
||||
async def consume(self, count = 1, timeout = 0.05, encoding = "utf-8"):
|
||||
|
||||
"""
|
||||
Get messages from the Kafka server.
|
||||
: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.
|
||||
"""
|
||||
|
||||
# Ensure connectivity to the server.
|
||||
# If not connected, return with failure immediately.
|
||||
if not await self.ensure_connection(): return []
|
||||
|
||||
# Make a variable that will hold the final results:
|
||||
messages = []
|
||||
|
||||
try:
|
||||
|
||||
# Read some messages:
|
||||
results = await self.__consumer.getmany(
|
||||
max_records = max(1, count),
|
||||
timeout_ms = int(timeout * 1_000)
|
||||
)
|
||||
|
||||
# Format the received messages:
|
||||
if results:
|
||||
for topic_partition, records in results.items():
|
||||
for record in records:
|
||||
messages.append(ConsumedKafkaMessage.from_aiokafka(
|
||||
message = record,
|
||||
deserializer = lambda x: self.__serializer.deserialize(x, encoding = encoding)
|
||||
))
|
||||
|
||||
# Debugging print if something went wrong:
|
||||
except Exception as exception: self.__printer(exception)
|
||||
|
||||
# Done here:
|
||||
return messages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BidirectionalKafka:
|
||||
|
||||
# The 'roles' that the instance can take.
|
||||
# The master talks on the channel (topic) that the slave listens on and vice versa.
|
||||
# Master-Slave is only for deciding who talks on which channel and who listens on which.
|
||||
# In a two-party system, one must be the master, the other must be the slave.
|
||||
# There are no extra privileges that the master enjoys. The naming convention was borrowed from common protocols
|
||||
# used in electronics (like I2C).
|
||||
ROLE_MASTER = 1
|
||||
ROLE_SLAVE = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
role,
|
||||
topic,
|
||||
ack_topic: str = None,
|
||||
group: str = None,
|
||||
serializer = None,
|
||||
debug = True,
|
||||
debug_prefix = "Kafka (B) | ",
|
||||
**kwargs
|
||||
):
|
||||
|
||||
"""
|
||||
Creates a walkie-talkie type setup to use Kafka in a bidirectional manner. Fo more information on all the
|
||||
individual methods, please read through the doc-strings of the component classes 'ProducerKafka', and
|
||||
'ConsumerKafka'.
|
||||
:param role: Select from "ROLE_MASTER" and "ROLE_SLAVE". Between the two parties that are talking, one will be
|
||||
the master and the other will be the slave. The channel that the master uses to speak will the one the slave
|
||||
uses to listen, and vice versa.
|
||||
:param topic: The topic to communicate on. Will be the same between the master and the slave.
|
||||
:param ack_topic: Explicitly provide this for the second channel, or it will be created from the name of the
|
||||
topic itself. Will be the same between the master and the slave.
|
||||
:param group: The group to assign the instance to.
|
||||
: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.
|
||||
"""
|
||||
|
||||
# Not down the basic variables:
|
||||
self.__role = role
|
||||
self.__topic = topic
|
||||
self.__ack_topic = ack_topic or topic + "Ack"
|
||||
self.__group = group
|
||||
|
||||
# In case the current instance is the master,
|
||||
# it will talk on "topic", and listen on "ack_topic":
|
||||
if self.__role == self.ROLE_MASTER:
|
||||
self.__producer_kwargs = kwargs.copy()
|
||||
self.__producer = ProducerKafka(
|
||||
topic = self.__topic,
|
||||
serializer = serializer,
|
||||
debug = debug,
|
||||
debug_prefix = debug_prefix.strip() + " (P) | ",
|
||||
**self.__producer_kwargs
|
||||
)
|
||||
self.__consumer_kwargs = kwargs.copy()
|
||||
self.__consumer_kwargs["group_id"] = self.__group
|
||||
self.__consumer = ConsumerKafka(
|
||||
topic = self.__ack_topic,
|
||||
group = group,
|
||||
serializer = serializer,
|
||||
debug = debug,
|
||||
debug_prefix = debug_prefix.strip() + " (C) | ",
|
||||
**self.__consumer_kwargs
|
||||
)
|
||||
|
||||
# On the other hand, if the current instance is a slave,
|
||||
# It will listen on "topic", and talk on "ack_topic":
|
||||
else:
|
||||
self.__producer_kwargs = kwargs.copy()
|
||||
self.__producer = ProducerKafka(
|
||||
topic = self.__ack_topic,
|
||||
serializer = serializer,
|
||||
debug = debug,
|
||||
debug_prefix = debug_prefix.strip() + " (P) | ",
|
||||
**self.__producer_kwargs
|
||||
)
|
||||
self.__consumer_kwargs = kwargs.copy()
|
||||
self.__consumer_kwargs["group_id"] = self.__group
|
||||
self.__consumer = ConsumerKafka(
|
||||
topic = self.__topic,
|
||||
group = group,
|
||||
serializer = serializer,
|
||||
debug = debug,
|
||||
debug_prefix = debug_prefix.strip() + " (C) | ",
|
||||
**self.__consumer_kwargs
|
||||
)
|
||||
|
||||
def enable_debug(self):
|
||||
self.__producer.enable_debug()
|
||||
self.__consumer.enable_debug()
|
||||
|
||||
def disable_debug(self):
|
||||
self.__producer.disable_debug()
|
||||
self.__consumer.disable_debug()
|
||||
|
||||
async def ensure_connection(self):
|
||||
await self.__producer.ensure_connection()
|
||||
await self.__consumer.ensure_connection()
|
||||
|
||||
async def close(self):
|
||||
await self.__producer.close()
|
||||
await self.__consumer.close()
|
||||
|
||||
async def produce(self, message, encoding = "utf-8"):
|
||||
return await self.__producer.produce(message, encoding = encoding)
|
||||
|
||||
async def consume(self, count = 1, timeout = 0.05, encoding = "utf-8"):
|
||||
return await self.__consumer.consume(count = count, timeout = timeout, encoding = encoding)
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** MAIN PROGRAM ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
# from data_models.kafka_message import KafkaMessage
|
||||
|
||||
ssl_ctx = get_ssl_context(
|
||||
ca_file = r"/home/developer/PycharmProjects/utils/cred/kafka/cert_authority.pem",
|
||||
cert_file = r"/home/developer/PycharmProjects/utils/cred/kafka/fullchain.pem",
|
||||
key_file = r"/home/developer/PycharmProjects/utils/cred/kafka/privkey.pem"
|
||||
)
|
||||
|
||||
async def consumer_test():
|
||||
|
||||
consumer = ConsumerKafka(
|
||||
topic = "tick-listners",
|
||||
# group_id = "assessImg",
|
||||
group_id = "updateMedia",
|
||||
bootstrap_servers = "wtt.ditscentre.in:9092",
|
||||
security_protocol = "SSL",
|
||||
ssl_context = ssl_ctx
|
||||
)
|
||||
await consumer.connect()
|
||||
await asyncio.sleep(1.5)
|
||||
print("READY!")
|
||||
|
||||
while True:
|
||||
messages = await consumer.consume(count = 1)
|
||||
for message in messages: print(message)
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
async def producer_test():
|
||||
|
||||
producer = ProducerKafka(
|
||||
topic = "tick-listners",
|
||||
bootstrap_servers = "del.ditscentre.in:9092",
|
||||
security_protocol = "SSL",
|
||||
ssl_context = ssl_ctx
|
||||
)
|
||||
await producer.connect()
|
||||
print("READY!")
|
||||
|
||||
while True:
|
||||
success = await producer.produce({"name": "Bhopli", "color": "orange"})
|
||||
print("PRODUCED:", success)
|
||||
await asyncio.sleep(0.25)
|
||||
|
||||
await producer.close()
|
||||
|
||||
|
||||
asyncio.run(consumer_test())
|
||||
@@ -0,0 +1,586 @@
|
||||
"""
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,257 @@
|
||||
"""
|
||||
|
||||
AUTHOR:
|
||||
|
||||
Khushal P Soonderji
|
||||
|
||||
DATE:
|
||||
|
||||
Monday, 30th Dec., 2024.
|
||||
|
||||
OBJECTIVE:
|
||||
|
||||
To provide a standardized structure for Kafka messages.
|
||||
|
||||
REFERENCES:
|
||||
|
||||
N/A
|
||||
|
||||
DOWNLOADS:
|
||||
|
||||
N/A
|
||||
|
||||
"""
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** IMPORT ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# To make sibling directories accessible for imports:
|
||||
import sys
|
||||
sys.path.append(".")
|
||||
sys.path.append("..")
|
||||
|
||||
# For making data behaviour_models:
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator, AwareDatetime
|
||||
from typing import Optional, Literal, Union, Dict, List, Any
|
||||
|
||||
# Related to Google:
|
||||
from google.auth.transport.requests import Request
|
||||
from google.oauth2.credentials import Credentials
|
||||
|
||||
# My utils:
|
||||
from utils_v2.string import json
|
||||
from utils_v2.string import regex
|
||||
from utils_v2.date_time import date_time
|
||||
|
||||
# To work with date and time:
|
||||
import datetime
|
||||
import dateparser
|
||||
|
||||
# To make API calls:
|
||||
import httpx
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** MACROS / ONE-TIME INIT ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# --- Nothing Yet
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** VARIABLES ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
# --- Nothing Yet
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** FUNCTIONS ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
class ConsumedKafkaMessage(BaseModel):
|
||||
|
||||
topic: str = Field(
|
||||
description = "the topic on which this message was received",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
partition: int = Field(
|
||||
description = "the partition in which this message was received",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
offset: int = Field(
|
||||
description = "the message's no. in the partition",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
headers: List[Any] = Field(
|
||||
description = "the headers received with the message",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
key: Any = Field(
|
||||
description = "the key with which this message is associated; important for partition management",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
value: Any = Field(
|
||||
description = "the actual payload of the message",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
ts: AwareDatetime | None = Field(
|
||||
description = "the time (utc) at which this message was sent to the queue",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
tsType: Literal[
|
||||
"createTime", # ...... The time at which the producer produced the message.
|
||||
"logAppendTime", # ... The time at which the message was received by the broker.
|
||||
None # ............... Unknown.
|
||||
] = Field(
|
||||
description = "to understand the source of the timestamp",
|
||||
frozen = True
|
||||
)
|
||||
|
||||
# ┏┓ ┏•
|
||||
# ┃ ┏┓┏┓╋┓┏┓
|
||||
# ┗┛┗┛┛┗┛┗┗┫
|
||||
# ┛
|
||||
|
||||
class Config:
|
||||
extra = "ignore"
|
||||
populate_by_name = True
|
||||
|
||||
# ┓┏ ┓• ┓ •
|
||||
# ┃┃┏┓┃┓┏┫┏┓╋┓┏┓┏┓
|
||||
# ┗┛┗┻┗┗┗┻┗┻┗┗┗┛┛┗
|
||||
|
||||
@field_validator("ts", mode = "before")
|
||||
def parse_dates(cls, value):
|
||||
if value is None: return None
|
||||
if not isinstance(value, datetime.datetime):
|
||||
parsed = date_time.parse_date_time(
|
||||
value,
|
||||
date_formats = ["%Y-%m-%d %H:%M:%S"],
|
||||
timezone = None
|
||||
)
|
||||
value = parsed if isinstance(parsed, datetime.datetime) else dateparser.parse(value)
|
||||
if isinstance(value, datetime.datetime): value = date_time.to_timezone(value, date_time.TIMEZONE_UTC)
|
||||
return value
|
||||
|
||||
# ┏┓ •
|
||||
# ┃┃┏┓┏┓┏┓┏┓┏┓╋┓┏┓┏
|
||||
# ┣┛┛ ┗┛┣┛┗ ┛ ┗┗┗ ┛
|
||||
# ┛
|
||||
|
||||
pass
|
||||
|
||||
# ┏┓ ┏┓
|
||||
# ┃ ┓┏┏╋┏┓┏┳┓ ┣ ┓┏┏┓┏┏
|
||||
# ┗┛┗┻┛┗┗┛┛┗┗ ┻ ┗┻┛┗┗┛
|
||||
|
||||
@staticmethod
|
||||
def from_aiokafka(
|
||||
message,
|
||||
deserializer = None
|
||||
):
|
||||
|
||||
"""
|
||||
To populate this model directly from the output of the 'aiokafka' library.
|
||||
:param message: The raw message from the library.
|
||||
:param deserializer: The function to use to deserialize the contents of the message.
|
||||
:return: The standardized Kafka consumed message.
|
||||
"""
|
||||
|
||||
# Extract the key and value:
|
||||
key = message.key
|
||||
value = message.value
|
||||
if deserializer:
|
||||
if key: key = deserializer(key)
|
||||
if value: value = deserializer(value)
|
||||
|
||||
# Build and return the model:
|
||||
return ConsumedKafkaMessage(
|
||||
topic = message.topic,
|
||||
partition = message.partition,
|
||||
offset = message.offset,
|
||||
headers = message.headers or [],
|
||||
key = key,
|
||||
value = value,
|
||||
ts = message.timestamp / 1000.0 if message.timestamp else None,
|
||||
tsType = {
|
||||
0: "createTime",
|
||||
1: "logAppendTime"
|
||||
}.get(message.timestamp_type)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_confluent_kafka(
|
||||
message,
|
||||
deserializer = None
|
||||
):
|
||||
|
||||
"""
|
||||
To populate this model directly from the output of the 'aiokafka' library.
|
||||
:param message: The raw message from the library.
|
||||
:param deserializer: The function to use to deserialize the contents of the message.
|
||||
:return: The standardized Kafka consumed message.
|
||||
"""
|
||||
|
||||
# If the message was null or an error:
|
||||
if message is None or message.error(): return None
|
||||
|
||||
# Figure out the timestamp:
|
||||
raw_ts = message.timestamp()
|
||||
ts_type = raw_ts[0] if raw_ts else None
|
||||
ts = raw_ts[1] / 1_000.0 if raw_ts else None
|
||||
|
||||
# Extract the key and value:
|
||||
key = message.key()
|
||||
value = message.value()
|
||||
if deserializer:
|
||||
if key: key = deserializer(key)
|
||||
if value: value = deserializer(value)
|
||||
|
||||
# Build and return the model:
|
||||
return ConsumedKafkaMessage(
|
||||
topic = message.topic(),
|
||||
partition = message.partition(),
|
||||
offset = message.offset(),
|
||||
headers = message.headers() or [],
|
||||
key = key,
|
||||
value = value,
|
||||
ts = ts,
|
||||
tsType = {
|
||||
1: "createTime",
|
||||
2: "logAppendTime"
|
||||
}.get(ts_type)
|
||||
)
|
||||
|
||||
|
||||
# *****************************************************************************************************************
|
||||
# ***** ****
|
||||
# *** MAIN PROGRAM ***
|
||||
# ***** ****
|
||||
# *****************************************************************************************************************
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user