Edits history of script submission #23952 for ' Publish message (kafka)'

• python3

  Ap­pro­ved version

  Edit in Windmill

  Embed

  # 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()}
  

  Submitted by hugo989 9 hours ago