Using python compiled protobuf pb2 as key and value serializer

Man, I was trying to do the same the past week. After struggling I finally got something working - not the best way - but it works well enough.

So initially I used this python compiler: https://github.com/danielgtaylor/python-betterproto to generate the *.py files with dataclasses / type hinting.

Then, I was able to create Faust.Record classes dynamically by using a helper:

import abc
import inspect
from typing import Type

import betterproto
import faust

GENERATED_SUFFIX = "__FaustRecord_Auto"


def _import_relative_class(module: str, klass_name: str):
    resolved_import = __import__(module, fromlist=[klass_name])
    klass = getattr(resolved_import, klass_name)
    return klass


def _is_record(attype: Type):
    return (
        inspect.isclass(attype)
        and isinstance(attype, betterproto.Message)
        or isinstance(attype, abc.ABCMeta)
    )


def _build_record_annotations(klass: Type):
    annotations = {}
    for atname, attype in klass.__annotations__.items():
        if _is_record(attype):
            annotations[atname] = make_faust_record(attype)
        elif isinstance(attype, str):
            subklass = _import_relative_class(klass.__module__, attype)
            annotations[atname] = make_faust_record(subklass)
        else:
            annotations[atname] = attype

    return annotations


def make_faust_record(klass: Type):
    type_name = f"{klass.__name__}{GENERATED_SUFFIX}"
    record_type = type(type_name, (faust.Record, klass), {})
    record_type.__annotations__ = _build_record_annotations(klass)
    record_type._init_subclass()

    return record_type

Now you can use it like:

import faust
from proto.your_models import YourModel # Import your generated proto here
from faust_converter import make_faust_record


app = faust.App(
    'faust-consumer',
    broker='kafka://',
    store="memory://",
    cache="memory://",
)

model_record = make_faust_record(YourModel)

topic = app.topic(
    'topic',
    value_type=model_record
)


@app.agent(topic)
async def consume(topic):
    async for event in topic:
        print(event)


if __name__ == "__main__":
    app.main()

I was also experimenting with using Protobuf with Faust.

Mentioned below is the solution using Faust Serialiser Codecs. faust-protobuf https://github.com/hemantkashniyal/faust-protobuf

proto_serializer.py

from faust.serializers import codecs
from typing import Any

from google.protobuf import json_format
from google.protobuf.json_format import MessageToJson
from google.protobuf.json_format import MessageToDict
from google.protobuf import text_format
from google.protobuf.text_format import MessageToString
from google.protobuf.text_format import MessageToBytes

class ProtobufSerializer(codecs.Codec):
    def __init__(self, pb_type: Any):
        self.pb_type = pb_type
        super(self.__class__, self).__init__()

    def _dumps(self, pb: Any) -> bytes:
        return pb.SerializeToString()

    def _loads(self, s: bytes) -> Any:
        pb = self.pb_type()
        pb.ParseFromString(s)
        return pb

app.py

import faust
from google.protobuf.json_format import MessageToJson

from .proto.greetings_pb2 import Greeting

from .proto_serializer import ProtobufSerializer

app = faust.App(
    'faust-consumer',
    broker='kafka://', # TODO: update kafka endpoint
    store="memory://",
    cache="memory://",
)

greetings_schema = faust.Schema(
    key_serializer=ProtobufSerializer(pb_type=Greeting),
    value_serializer=ProtobufSerializer(pb_type=Greeting),
)

topic = app.topic(
    'greetings',
    schema=greetings_schema
)

@app.agent(topic)
async def consume(topic):
    async for event in topic:
        print(MessageToJson(event))

@app.timer(5)
async def produce():
    for i in range(10):
        data = Greeting(hello="world", message=i)
        await consume.send(value=data)

if __name__ == "__main__":
    app.main()

I was able to do it by creating a Serializer class as so:

import faust
from abc import ABCMeta, abstractmethod
from google.protobuf.json_format import MessageToDict
from faust.serializers.codecs import Codec
from importlib import import_module


def get_proto(topic_name, only_pk=False):
    if not hasattr(get_proto, "topics"):
        setattr(get_proto, "topics", dict())

    get_proto.topics[topic_name] = import_module(
        "protodef.{}_pb2".format(topic_name)
    ).__getattribute__(topic_name.split(".")[-1])

    if only_pk:
        return getattr(get_proto, "topics").get(topic_name).PK
    else:
        return getattr(get_proto, "topics").get(topic_name)


class ProtoSerializer(Codec, metaclass=ABCMeta):
    @abstractmethod
    def only_key(self):
        ...

    def as_proto(self, topic_name):
        self._proto = get_proto(topic_name, self.only_key())
        return self

    def _loads(self, b):
        data = MessageToDict(
            self._proto.FromString(b),
            preserving_proto_field_name=True,
            including_default_value_fields=True,
        )

        # remove the key object from the unserialized message
        data.pop("key", None)
        return data

    def _dumps(self, o):
        # for deletes
        if not o:
            return None

        obj = self._proto()

        # add the key object to them message before serializing
        if hasattr(obj, "PK"):
            for k in obj.PK.DESCRIPTOR.fields_by_name.keys():
                if k not in o:
                    raise Exception(
                        "Invalid object `{}` for proto `{}`".format(o, self._proto)
                    )
                setattr(obj.key, k, o[k])

        for k, v in o.items():
            if hasattr(obj, k):
                setattr(obj, k, v)
            else:
                ghost.debug(
                    "Invalid value-attribute `%s` for proto `%s`", k, self._proto
                )

        return obj.SerializeToString()


class ProtoValue(ProtoSerializer):
    def only_key(self):
        return False


class ProtoKey(ProtoSerializer):
    def only_key(self):
        return True

and then use it as follows:

import faust
from utils.serializer import ProtoKey, ProtoValue


app = faust.App(
    'faust-consumer',
    broker='kafka://',
    store="memory://",
    cache="memory://",
)


topic = app.topic(
    'topic',
    key_serializer=ProtoKey().as_proto('topic'),
    value_serializer=ProtoValue().as_proto('topic')
)


@app.agent(topic)
async def consume(topic):
    async for event in topic:
        print(event)


if __name__ == "__main__":
    app.main()