Function _kafka_chunk

misc/python/materialize/workload_replay/data.py:94–175

Generate random data and produce to Kafka.

(
    kafka_port: int,
    sr_port: int,
    topic: str,
    debezium: bool,
    column_dicts: list[dict[str, Any]],
    num_rows: int,
    rng_seed: int,
)

def _kafka_chunk(
    kafka_port: int,
    sr_port: int,
    topic: str,
    debezium: bool,
    column_dicts: list[dict[str, Any]],
    num_rows: int,
    rng_seed: int,
) -> int:
    """Generate random data and produce to Kafka."""
    rng = random.Random(rng_seed)
    columns = [
        Column(c["name"], c["type"], c["nullable"], c["default"], c.get("data_shape"))
        for c in column_dicts
    ]

    producer, serializer, key_serializer, sctx, ksctx, col_names = get_kafka_objects(
        topic,
        tuple(columns),
        debezium,
        sr_port,
        kafka_port,
    )

    now_ms = int(time.time() * 1000)
    source_struct: dict[str, Any] | None = None
    if debezium:
        source_struct = {
            "version": "0",
            "connector": "mysql",
            "name": "materialize-generator",
            "ts_ms": now_ms,
            "snapshot": None,
            "db": "db",
            "sequence": None,
            "table": topic.split(".")[-1],
            "server_id": 0,
            "gtid": None,
            "file": "binlog.000001",
            "pos": 0,
            "row": 0,
            "thread": None,
            "query": None,
        }

    producer.poll(0)
    for _ in range(num_rows):
        row = [col.kafka_value(rng) for col in columns]
        while True:
            try:
                if debezium:
                    after_value = dict(zip(col_names, row))
                    envelope_value = {
                        "before": None,
                        "after": after_value,
                        "source": source_struct,
                        "op": "c",
                        "ts_ms": now_ms,