collect.py

#!/usr/bin/env python
# coding: utf8

"""plot offset"""

import os
import concurrent.futures
from datetime import datetime, timedelta
import subprocess
import json
import logging
import redis
from pydantic import BaseModel
import dotenv

import pandas as pd
import matplotlib.pyplot as plt

dotenv.load_dotenv()

REDIS_OFFSET_CALC_PREFIX = 'offset-calc-'
REDIS_KEY_TTL = 60 * 60


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)

TOPIC1 = os.getenv('TOPIC1')
TOPIC2 = os.getenv('TOPIC2')
if not TOPIC1 or not TOPIC2:
    raise ValueError("Environment variables TOPIC1 and TOPIC2 must be set.")

REDIS_URL = os.getenv('REDIS_URL')
if not REDIS_URL:
    raise ValueError("Environment variable REDIS_URL must be set.")

redis_client = redis.Redis.from_url(REDIS_URL)
redis_client.ping()

current_user = os.getenv('USER', 'ec2-user')
if not current_user:
    raise ValueError("Environment variable USER must be set.")

KCAT = f'docker run --rm -v /home/{current_user}/.config/kcat.conf:/root/.config/kcat.conf edenhill/kcat:1.7.1'.split(' ')

def kcat_meta_listing(topic: str | None = None) -> None | dict:
    command = [*KCAT, '-L', '-J']
    if topic:
        command.extend(['-t', topic])
    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT)
        result = json.loads(output)
        return result
    except subprocess.CalledProcessError as e:
        print(f"Error running kcat: {e.output.decode()}")
        return None

def get_brokers():
    result = kcat_meta_listing()
    if result is None:
        return None
    return result['brokers']

def get_topic_partitions(topic: str) -> list | None:
    result = kcat_meta_listing(topic)
    if result is None:
        return None
    return result['topics'][0]['partitions']

def kcat_query_offset(dt: datetime | str | int | float, topic: str, partition: int = 0, datetime_format='%Y-%m-%d %H:%M:%S') -> int:
    """
    dt: %Y-%m-%d %H:%M:%S, PT timezone by default.
    Query options (-Q):
        -t <t>:<p>:<ts> Get offset for topic <t>, partition <p>, timestamp <ts>.
        Timestamp is the number of milliseconds since epoch UTC.
        Requires broker >= 0.10.0.0 and librdkafka >= 0.9.3.
        Multiple -t .. are allowed but a partition must only occur once.
    """
    match dt:
        case datetime():
            dt_timestamp = int(dt.timestamp() * 1000)
        case str():
            dt_timestamp = int(datetime.strptime(dt, datetime_format).timestamp() * 1000)
        case int() | float():
            dt_timestamp = int(dt)
        case _:
            raise ValueError(f"Unsupported type for dt: {type(dt)}")
    log.info(f"Querying offset for topic {topic}, partition {partition}, timestamp {dt_timestamp} ({dt})")

    command = [*KCAT, '-Q', '-J', '-t', f'{topic}:{partition}:{dt_timestamp}']
    output = subprocess.check_output(command, stderr=subprocess.STDOUT)
    result = json.loads(output)
    return result[topic]['0']['offset']

"""
Consumer options:
  -o <offset>        Offset to start consuming from:
                     beginning | end | stored |
                     <value>  (absolute offset) |
                     -<value> (relative offset from end)
                     s@<value> (timestamp in ms to start at)
                     e@<value> (timestamp in ms to stop at (not included))
  -e                 Exit successfully when last message received
  -f <fmt..>         Output formatting string, see below.
                     Takes precedence over -D and -K.
  -J                 Output with JSON envelope
  -s key=<serdes>    Deserialize non-NULL keys using <serdes>.
  -s value=<serdes>  Deserialize non-NULL values using <serdes>.
  -s <serdes>        Deserialize non-NULL keys and values using <serdes>.
                     Available deserializers (<serdes>):
                       <pack-str> - A combination of:
                                    <: little-endian,
                                    >: big-endian (recommended),
                                    b: signed 8-bit integer
                                    B: unsigned 8-bit integer
                                    h: signed 16-bit integer
                                    H: unsigned 16-bit integer
                                    i: signed 32-bit integer
                                    I: unsigned 32-bit integer
                                    q: signed 64-bit integer
                                    Q: unsigned 64-bit integer
                                    c: ASCII character
                                    s: remaining data is string
                                    $: match end-of-input (no more bytes remaining or a parse error is raised).
                                       Not including this token skips any
                                       remaining data after the pack-str is
                                       exhausted.
                       avro       - Avro-formatted with schema in Schema-Registry (requires -r)
                     E.g.: -s key=i -s value=avro - key is 32-bit integer, value is Avro.
                       or: -s avro - both key and value are Avro-serialized
  -r <url>           Schema registry URL (when avro deserializer is used with -s)
  -D <delim>         Delimiter to separate messages on output
  -K <delim>         Print message keys prefixing the message
                     with specified delimiter.
  -O                 Print message offset using -K delimiter
  -c <cnt>           Exit after consuming this number of messages
  -Z                 Print NULL values and keys as "NULL" instead of empty.
                     For JSON (-J) the nullstr is always null.
  -u                 Unbuffered output
"""


def kcat_consume_by_time(start_time: str, topic: str, end_time: str | None = None,
                         time_delta: timedelta = timedelta(minutes=5), partition: int | None = None, timeout: int = 10) -> None | str:
    """
    start_time: %Y-%m-%d %H:%M:%S, PT timezone by default.
    Consume messages from a topic starting from a specific time.
    """
    start_timestamp, end_timestamp = time_range(start_time, end_time, time_delta)
    log.info(f"Consuming messages from topic {topic} starting at {start_time} ({start_timestamp}) to {end_time} ({end_timestamp})")

    command = ['-C', '-t', topic, '-p', str(partition), '-o', f's@{start_timestamp}', '-o', f'e@{end_timestamp}', '-J', '-e']
    log.info(f"Running command: {' '.join(command)}")
    try:
        output = subprocess.check_output(KCAT + command, stderr=subprocess.STDOUT, timeout=timeout, text=True)
        return output
        # result = json.loads(output)
        # return result
    except subprocess.TimeoutExpired as e:
        log.error(f"Timeout expired: {e}")
        return None
    except subprocess.CalledProcessError as e:
        log.error(f"Error running kcat: {e.output.decode()}")
        return None

def kcat_consume_latest(topic: str, partition: int = 0, timeout: int = 10) -> None | list:
    """
    Consume the latest messages from a topic.
    """
    command = [*KCAT, '-C', '-t', topic, '-p', str(partition), '-o', 'end', '-e', '-J', '-T']
    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT, timeout=timeout)
        result = json.loads(output)
        return result
    except subprocess.CalledProcessError as e:
        print(f"Error running kcat: {e.output.decode()}")
        return None

def time_range(start_time: str, end_time: str | None = None, time_delta: timedelta | None = None, timestamp: bool=True) -> tuple[datetime, datetime] | tuple[float, float]:
    """
    Generate a list of timestamps in milliseconds from start_time to end_time with a given time_delta.
    start_time: %Y-%m-%d %H:%M:%S, PT timezone by default.
    """
    start_dt = datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S")
    if end_time:
        end_dt = datetime.strptime(end_time, "%Y-%m-%d %H:%M:%S")
    else:
        log.info('end_time is not specified, use start_time + time_delta (%s)', time_delta if time_delta else f'fallback: {timedelta(minutes=5)}')
        end_dt = start_dt + (time_delta or timedelta(minutes=5))
    return (start_dt.timestamp() * 1000, end_dt.timestamp() * 1000) if timestamp else (start_dt, end_dt)

class OffsetData(BaseModel):
    time: str
    start_offset: int
    end_offset: int
    offset_diff: int

def per_minute_offset(start_time: str, topic: str, partition: int = 0, minutes: int=30, redis_key_prefix: str=REDIS_OFFSET_CALC_PREFIX, redis_key_ttl: int=REDIS_KEY_TTL) -> list:
    """
    Collect offsets for each minute starting from start_time for a given topic and partition.
    """
    log.info(f"Collecting offsets for topic {topic} from {start_time} every minute.")
    time_delta = timedelta(minutes=1)
    offsets = []

    for i in range(minutes):
        current_time = (datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S") + timedelta(minutes=i)).strftime("%Y-%m-%d %H:%M:%S")

        start_dt, end_dt = time_range(current_time, time_delta=time_delta, timestamp=False)
        start_offset = kcat_query_offset(start_dt, topic, partition=partition)
        end_offset = kcat_query_offset(end_dt, topic, partition=partition)
        offsets.append(OffsetData(time=current_time, start_offset=start_offset, end_offset=end_offset, offset_diff=end_offset - start_offset))
        log.info(f"Offset at {current_time}: {offsets[-1]}")

    global redis_client
    key = f"{redis_key_prefix}{topic}-{partition}-{start_time.replace(' ', 'T')}-{minutes}min"
    redis_client.setex(key, redis_key_ttl, json.dumps([offset.model_dump() for offset in offsets]))
    log.info(f"Offsets saved to Redis with key: {key}")

    return offsets


def collect_offsets_in_parallel(start_date: str, topic: str, partition: int = 0, total_minutes: int = 60, chunk_size: int = 10, dt_format='%Y-%m-%d %H:%M:%S') -> None:
    """
    Collect offsets in parallel using multithreading.
    """
    with concurrent.futures.ProcessPoolExecutor() as executor:
        futures = []
        for i in range(0, total_minutes, chunk_size):
            start_time = (datetime.strptime(start_date, dt_format) + timedelta(minutes=i)).strftime(dt_format)
            log.info('start_time: %s', start_time)
            futures.append(executor.submit(per_minute_offset, start_time, topic, partition, minutes=chunk_size))

        for future in concurrent.futures.as_completed(futures):
            try:
                data = future.result()
                log.info("Offsets collected: \n%s", json.dumps([offset.model_dump() for offset in data], indent=2))
            except Exception as e:
                log.error(f"Error occurred: {e}")


def collect_offset_data_from_redis():
    global redis_client

    redis_keys = redis_client.keys(f"{REDIS_OFFSET_CALC_PREFIX}*")
    offset_data = []
    for key in redis_keys:
        data = redis_client.get(key)
        if data:
            offset_data.extend(json.loads(data.decode('utf-8')))
            log.info(f"Loaded {len(offset_data)} records from Redis key: {key.decode('utf-8')}")
    return sorted(offset_data, key=lambda x: x['time'])


def plot_offset(topic: str = TOPIC1) -> None:
    """
    Plot offsets over time for a specific topic.
    """
    offset_data = collect_offset_data_from_redis()
    offsets_df = pd.DataFrame.from_records(offset_data)
    offsets_df['time'] = pd.to_datetime(offsets_df['time'])

    plt.figure(figsize=(12, 6))
    plt.plot(offsets_df['time'], offsets_df['offset_diff'], marker='o')
    plt.title(f"Data counts for {topic} over time (per minute)")
    plt.xlabel('Time')
    plt.ylabel('Offset')
    plt.xticks(rotation=45)
    plt.grid()
    plt.tight_layout()
    plt.savefig('offsets_over_time.png')
    plt.show()

requirements.txt

kafka-python==2.2.10
python-dotenv
confluent-kafka
pandas
polars
numpy
scikit-learn
matplotlib
seaborn
jupyterlab
duckdb
pydantic
redis