Publish message

# py312

# requirements:

# confluent-kafka==2.16.1

import json

import time

from typing import Optional, TypedDict

from confluent_kafka import KafkaError, KafkaException, Producer

class kafka(TypedDict):

    brokers: list

    security: dict

PROTOCOLS = {

    "PLAINTEXT": "plaintext",

    "SASL_PLAINTEXT": "sasl_plaintext",

    "SSL": "ssl",

    "SASL_SSL": "sasl_ssl",

    "SASL_SSL_OAUTHBEARER": "sasl_ssl",

}

def producer_config(auth: kafka) -> dict:

    # Same librdkafka settings as Windmill's Kafka trigger reads from this resource.

    security = auth.get("security") or {}

    label = security.get("label") or "PLAINTEXT"

    if "GSSAPI" in label:

        raise Exception(

            f"{label} is not supported: the Kafka client this script uses is built without Kerberos (GSSAPI)"

        )

    if label not in PROTOCOLS:

        raise Exception(f"Unknown security label {label!r}: expected one of {', '.join(PROTOCOLS)}")

    conf = {

        "bootstrap.servers": ",".join(auth["brokers"]),

        "security.protocol": PROTOCOLS[label],

        "message.timeout.ms": 30000,

    }

    if label in ("SASL_PLAINTEXT", "SASL_SSL"):

        conf["sasl.mechanism"] = security["mechanism"]

        conf["sasl.username"] = security["username"]

        conf["sasl.password"] = security["password"]

    if label == "SASL_SSL_OAUTHBEARER":

        conf["sasl.mechanism"] = "OAUTHBEARER"

        conf["sasl.oauthbearer.method"] = "oidc"

        conf["sasl.oauthbearer.client.id"] = security["client_id"]

        conf["sasl.oauthbearer.client.secret"] = security["client_secret"]

        conf["sasl.oauthbearer.token.endpoint.url"] = security["token_endpoint_url"]

        conf["https.ca.location"] = "probe"

        if security.get("scope"):

            conf["sasl.oauthbearer.scope"] = security["scope"]

        if security.get("extensions"):

            conf["sasl.oauthbearer.extensions"] = security["extensions"]

    for field, key in (

        ("ca", "ssl.ca.pem"),

        ("certificate", "ssl.certificate.pem"),

        ("key", "ssl.key.pem"),

        ("key_password", "ssl.key.password"),

    ):

        if security.get(field):

            conf[key] = security[field]

    return conf

def main(

    auth: kafka,

    topic: str,

    payload,

    key: Optional[str] = None,

    headers: Optional[dict] = None,

):

    # None is a tombstone (a null value, which deletes the key on a compacted topic).

    if payload is None:

        value = None

    else:

        value = (payload if isinstance(payload, str) else json.dumps(payload)).encode()

    errors = []  # client-level errors (bad credentials, TLS, unreachable brokers)

    producer = Producer({**producer_config(auth), "error_cb": errors.append})

    reports = []

    producer.produce(

        topic,

        value=value,

        key=str(key).encode() if key not in (None, "") else None,

        headers=[

            (k, (v if isinstance(v, str) else json.dumps(v)).encode())

            for k, v in (headers or {}).items()

        ]

        or None,

        on_delivery=lambda err, msg: reports.append((err, msg)),

    )

    # librdkafka retries refused connections until message.timeout.ms. Stop early only

    # on a fatal error, or when every broker is down and one refused the credentials:

    # a broker that is briefly unreachable, or one bad broker among several, may still

    # let the message through.

    deadline = time.monotonic() + 40

    while not reports and time.monotonic() < deadline:

        producer.poll(0.2)

        for e in errors:

            if e.fatal():

                raise KafkaException(e)

        auth_errors = [e for e in errors if e.code() == KafkaError._AUTHENTICATION]

        if auth_errors and any(e.code() == KafkaError._ALL_BROKERS_DOWN for e in errors):

            raise KafkaException(auth_errors[0])

    if not reports:

        raise Exception("Timed out waiting for the broker to acknowledge the message")

    err, msg = reports[0]

    if err is not None:

        # A timed-out delivery says nothing of why: add the last connection error.

        cause = [e for e in errors if e.code() != KafkaError._ALL_BROKERS_DOWN]

        raise Exception(err.str() + (f" (last error: {cause[-1].str()})" if cause else ""))

    return {"topic": msg.topic(), "partition": msg.partition(), "offset": msg.offset()}