from pyflink.common import Row
from pyflink.common import Types
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.datastream.connectors.kafka import (
    KafkaSink,
    KafkaRecordSerializationSchema,
    DeliveryGuarantee,
)
from pyflink.datastream.formats.json import JsonRowSerializationSchema

env = StreamExecutionEnvironment.get_execution_environment()
jar_paths = (
    "<path-to>/flink-json-1.17.1.jar",
    "<path-to>/flink-sql-connector-kafka-1.17.1.jar",
)
env.add_jars(*jar_paths)


def to_row(t: tuple):
    return Row(a=t[0], b=t[1])


data_stream = env.from_collection(collection=[(1, "2"), (3, "4")])

row_type_info = Types.ROW_NAMED(["a", "b"], [Types.INT(), Types.STRING()])
json_format = JsonRowSerializationSchema.builder().with_type_info(row_type_info).build()

sink = (
    KafkaSink.builder()
    .set_bootstrap_servers("localhost:29092")
    .set_record_serializer(
        KafkaRecordSerializationSchema.builder()
        .set_topic("test")
        .set_value_serialization_schema(json_format)
        .build()
    )
    .set_delivery_guarantee(DeliveryGuarantee.AT_LEAST_ONCE)
    .build()
)

## data_stream.map(to_row).print() # print works fine
data_stream.map(to_row).sink_to(sink)

env.execute()