Streaming en Tiempo Real con Pub/Sub, Dataflow y BigQuery usando Python

import json
import apache_beam as beam

from apache_beam.options.pipeline_options import (
    PipelineOptions,
    StandardOptions,
)


class ParseTelemetry(beam.DoFn):
    """Convierte el mensaje Pub/Sub en un registro analítico."""

    def process(self, message):
        try:
            payload = json.loads(message.decode("utf-8"))

            yield {
                "device_id": str(payload["device_id"]),
                "event_ts": str(payload["event_ts"]),
                "temperature": float(payload["temperature"]),
                "battery": float(payload.get("battery", 0)),
                "region": str(payload.get("region", "unknown")),
            }

        except (ValueError, KeyError, TypeError) as exc:
            # En producción, enviaríamos el evento a una DLQ
            # con suficiente contexto para poder reprocesarlo.
            yield beam.pvalue.TaggedOutput(
                "dead_letter",
                {
                    "raw_message": message.decode("utf-8", errors="replace"),
                    "error": str(exc),
                },
            )


class AddDerivedMetrics(beam.DoFn):
    """Añade métricas derivadas sin acoplarlas al origen."""

    def process(self, row):
        temperature = row["temperature"]

        row["temperature_alert"] = (
            temperature >= 80.0 or temperature <= -10.0
        )

        yield row


def run():
    options = PipelineOptions(
        streaming=True,
        save_main_session=True,
    )

    options.view_as(StandardOptions).streaming = True

    input_subscription = (
        "projects/PROJECT_ID/subscriptions/telemetry-sub"
    )

    output_table = (
        "PROJECT_ID:analytics.telemetry"
    )

    dead_letter_table = (
        "PROJECT_ID:analytics.telemetry_dead_letter"
    )

    with beam.Pipeline(options=options) as pipeline:

        parsed = (
            pipeline
            | "ReadFromPubSub"
            >> beam.io.ReadFromPubSub(
                subscription=input_subscription
            )
            | "ParseTelemetry"
            >> beam.ParDo(ParseTelemetry()).with_outputs(
                "dead_letter",
                main="valid",
            )
        )

        valid_rows = (
            parsed.valid
            | "AddDerivedMetrics"
            >> beam.ParDo(AddDerivedMetrics())
        )

        (
            valid_rows
            | "WriteToBigQuery"
            >> beam.io.WriteToBigQuery(
                output_table,
                schema={
                    "fields": [
                        {"name": "device_id", "type": "STRING"},
                        {"name": "event_ts", "type": "TIMESTAMP"},
                        {"name": "temperature", "type": "FLOAT64"},
                        {"name": "battery", "type": "FLOAT64"},
                        {"name": "region", "type": "STRING"},
                        {"name": "temperature_alert", "type": "BOOL"},
                    ]
                },
                write_disposition=(
                    beam.io.BigQueryDisposition.WRITE_APPEND
                ),
                create_disposition=(
                    beam.io.BigQueryDisposition.CREATE_IF_NEEDED
                ),
            )
        )

        (
            parsed.dead_letter
            | "WriteDeadLetter"
            >> beam.io.WriteToBigQuery(
                dead_letter_table,
                write_disposition=(
                    beam.io.BigQueryDisposition.WRITE_APPEND
                ),
                create_disposition=(
                    beam.io.BigQueryDisposition.CREATE_IF_NEEDED
                ),
            )
        )


if __name__ == "__main__":
    run()