diff --git a/runtime/transformations/components/data-sinks/hetida_platform_kafka_010_fb307c92-26d1-4f34-867e-ee5bf4e94479.py b/runtime/transformations/components/data-sinks/hetida_platform_kafka_010_fb307c92-26d1-4f34-867e-ee5bf4e94479.py new file mode 100644 index 000000000..2ec218c52 --- /dev/null +++ b/runtime/transformations/components/data-sinks/hetida_platform_kafka_010_fb307c92-26d1-4f34-867e-ee5bf4e94479.py @@ -0,0 +1,261 @@ +import asyncio +import datetime +import json +import logging +import math +from typing import Annotated + +import aiokafka +from pydantic import ( + AwareDatetime, + BaseModel, + Field, + ValidationInfo, + field_serializer, + field_validator, +) +from pydantic_settings import BaseSettings, NoDecode + +logger = logging.getLogger(__name__) + + +class KafkaConfig(BaseSettings): + """Class that loads environment variables regarding the kafka connection configuration. + + Attributes + ---------- + servers : list[str] + Kafka bootstrap servers, loaded from "HETIDA_PLATFORM_KAFKA_SERVERS". + At least one entry is required. + tenants : list[str] + Tenant ids the data is sent to, loaded from "HETIDA_PLATFORM_KAFKA_TENANTS". + At least one entry is required. + topic : str + Kafka topic, loaded from "HETIDA_PLATFORM_KAFKA_TIMESERIES_INGESTION_TOPIC". + Must not be empty. + + `servers` and `tenants` accept a JSON list, e.g. + HETIDA_PLATFORM_KAFKA_SERVERS='["kafka-1:9092","kafka-2:9092"]', or a plain + string, e.g. HETIDA_PLATFORM_KAFKA_SERVERS='kafka:9092', both from the + environment and when passed directly as argument. Entries are additionally + split at commas, e.g. HETIDA_PLATFORM_KAFKA_SERVERS='kafka-1:9092,kafka-2:9092'. + If "HETIDA_PLATFORM_KAFKA_SERVERS" or + "HETIDA_PLATFORM_KAFKA_TIMESERIES_INGESTION_TOPIC" is missing, creating the + config fails with a ValidationError naming the variable. + """ + + # NoDecode: pydantic-settings passes the raw environment value to the validator + # below instead of decoding it as JSON itself. + servers: Annotated[list[str], NoDecode] = Field( + alias="HETIDA_PLATFORM_KAFKA_SERVERS", + min_length=1, + description="Kafka bootstrap servers", + ) + tenants: Annotated[list[str], NoDecode] = Field( + alias="HETIDA_PLATFORM_KAFKA_TENANTS", + min_length=1, + description="Tenant ids the data is sent to", + ) + topic: str = Field( + alias="HETIDA_PLATFORM_KAFKA_TIMESERIES_INGESTION_TOPIC", + min_length=1, + description="Kafka topic", + ) + + @field_validator("servers", "tenants", mode="before") + @classmethod + def parse_json_string(cls, v): + # Handles values from the environment and passed as argument (e.g. the + # tenant input in main) alike. A string not starting with "[" is taken + # as a single entry. + if not isinstance(v, str): + return v + v = v.strip() + if v.startswith("["): + return json.loads(v) + return [v] if v else [] + + @field_validator("servers", "tenants") + @classmethod + def split_comma_separated(cls, v, info: ValidationInfo): + # One server / tenant per list entry, aiokafka does not split list entries. + # Runs after the list validation, so min_length does not cover entries + # consisting only of commas. + items = [item.strip() for entry in v for item in entry.split(",")] + items = [item for item in items if item] + if not items: + raise ValueError(f"At least one entry in {info.field_name} is required") + return items + + +class KafkaTemplate(BaseModel): + """Class that corresponds to the format of a single data point in kafka payloads + for the platform. + + Only numeric values are supported, since the platform stores time series values + as floating point numbers. Values that cannot be converted to float, e.g. text, + make the execution fail with a ValidationError. Data points with NaN or + infinite value are skipped by `main`. + + Attributes + ---------- + timestamp : AwareDatetime + Timezone aware datetime, serialized via `isoformat()`. + metric : str + Key of timeseries (=alias) in database. + value : float + Measured value. + """ + + timestamp: AwareDatetime + metric: str + value: float + + @field_serializer("timestamp") + def serialize_dt(self, timestamp: datetime.datetime, _info): + return timestamp.isoformat() + + +class KafkaHandler: + """Class that handles the delivery of messages to kafka. + + Use it as async context manager: the producer is started on entering the + `async with` block and stopped on leaving it. + + Attributes + ---------- + config : KafkaConfig + Kafka configuration info. + kafka_producer : aiokafka.AIOKafkaProducer | None + Started Kafka producer to send messages with. + None outside of the `async with` block. + max_data_points_per_message : int + Upper limit of data points per kafka message. Keeps messages well below the + default maximum message size of 1 MB, since every data point is contained + twice (in "payloads" and in "originalPayload"). + """ + + max_data_points_per_message = 1000 + + def __init__(self, config: KafkaConfig): + self.config = config + self.kafka_producer: aiokafka.AIOKafkaProducer | None = None + + async def __aenter__(self): + producer = aiokafka.AIOKafkaProducer(bootstrap_servers=self.config.servers) + try: + await producer.start() + except Exception: + # A failed start() leaves the client's connections open, and + # AIOKafkaProducer.__aexit__ is not called in that case. + await producer.stop() + raise + self.kafka_producer = producer + return self + + async def __aexit__(self, *exc_info): + if self.kafka_producer is None: + return + await self.kafka_producer.stop() + self.kafka_producer = None + + async def send_msg_to_kafka(self, payload: list[dict]): + """Send data points to kafka, once per configured tenant. + + The data points are split into messages of at most + `max_data_points_per_message` data points each. + + Parameters + ---------- + payload : list[dict] + Data points as dumped `KafkaTemplate` objects. + + Raises + ------ + RuntimeError + If called outside of the `async with` block. + ValueError + If a data point has a NaN or infinite value, which is not valid JSON. + """ + if self.kafka_producer is None: + raise RuntimeError( + "Kafka producer is not started - use KafkaHandler with `async with`." + ) + + deliveries = [] + for start in range(0, len(payload), self.max_data_points_per_message): + chunk = payload[start : start + self.max_data_points_per_message] + original_payload = json.dumps(chunk, allow_nan=False) + for tenant in self.config.tenants: + kafka_payload = { + "tenantId": tenant, + "topic": self.config.topic, + "origin": "UNKNOWN", + "originalPayload": original_payload, + "payloads": chunk, + } + deliveries.append( + await self.kafka_producer.send( + self.config.topic, + json.dumps(kafka_payload, allow_nan=False).encode(), + ) + ) + # send() only enqueues; wait for all acknowledgements together so aiokafka can batch. + await asyncio.gather(*deliveries) + logger.info( + "Sent %d data point(s) in %d message(s) to topic %s", + len(payload), + len(deliveries), + self.config.topic, + ) + + +# %% +# ***** DO NOT EDIT LINES BELOW ***** +# These lines may be overwritten if component details or inputs/outputs change. +COMPONENT_INFO = { + "inputs": { + "data": {"data_type": "MULTITSFRAME"}, + "tenant": {"data_type": "STRING"}, + }, + "outputs": {}, + "name": "Hetida Platform Kafka", + "category": "Data Sinks", + "description": "New created component", + "version_tag": "0.1.0", + "id": "fb307c92-26d1-4f34-867e-ee5bf4e94479", + "revision_group_id": "7e2c3659-ba5c-4046-9865-78764282a006", + "state": "RELEASED", + "released_timestamp": "2026-09-09T11:26:28.406786+00:00", +} + +from hdutils import parse_default_value # noqa: E402, F401 + + +async def main(*, data, tenant): + # entrypoint function for this component + # ***** DO NOT EDIT LINES ABOVE ***** + + # write your function code here. + config = KafkaConfig(HETIDA_PLATFORM_KAFKA_TENANTS=tenant) + + data_points = [ + KafkaTemplate(**record).model_dump() for record in data.to_dict(orient="records") + ] + # NaN and infinite values are not valid JSON, the platform would reject the whole message. + finite_data_points = [point for point in data_points if math.isfinite(point["value"])] + if len(finite_data_points) < len(data_points): + logger.warning( + "Skipping %d data point(s) with NaN or infinite value", + len(data_points) - len(finite_data_points), + ) + + if not finite_data_points: + logger.info("No data points to send") + return + + async with KafkaHandler(config) as kafka_handler: + await kafka_handler.send_msg_to_kafka(finite_data_points) + + +# %%