common/redis/serializers.py

# SPDX-FileCopyrightText: 2020 The Magma Authors.
# SPDX-FileCopyrightText: 2022 Open Networking Foundation <support@opennetworking.org>
#
# SPDX-License-Identifier: BSD-3-Clause

from typing import Callable, Generic, Type, TypeVar

import jsonpickle
from orc8r.protos.redis_pb2 import RedisState

T = TypeVar('T')


class RedisSerde(Generic[T]):
    """
    typeval (str): str representing the type of object the serde can
        de/serialize
    serializer (function (T, int) -> str):
                function called to serialize a value
    deserializer (function (str) -> T):
                function called to deserialize a value
    """

    def __init__(
        self,
        redis_type: str,
        serializer: Callable[[T, int], str],
        deserializer: Callable[[str], T],
    ):
        self.redis_type = redis_type
        self.serializer = serializer
        self.deserializer = deserializer

    def serialize(self, msg: T, version: int = 1) -> str:
        return self.serializer(msg, version)

    def deserialize(self, serialized_obj: str) -> T:
        return self.deserializer(serialized_obj)


def get_proto_serializer() -> Callable[[T, int], str]:
    """
    Return a proto serializer that serializes the proto, adds the associated
    version, and then serializes the RedisState proto to a string
    """
    def _serialize_proto(proto: T, version: int) -> str:
        serialized_proto = proto.SerializeToString()
        redis_state = RedisState(
            serialized_msg=serialized_proto,
            version=version,
            is_garbage=False,
        )
        return redis_state.SerializeToString()
    return _serialize_proto


def get_proto_deserializer(proto_class: Type[T]) -> Callable[[str], T]:
    """
    Return a proto deserializer that takes in a proto type to deserialize
    the serialized msg stored in the RedisState proto
    """
    def _deserialize_proto(serialized_rule: str) -> T:
        proto_wrapper = RedisState()
        proto_wrapper.ParseFromString(serialized_rule)
        serialized_proto = proto_wrapper.serialized_msg
        proto = proto_class()
        proto.ParseFromString(serialized_proto)
        return proto
    return _deserialize_proto


def get_json_serializer() -> Callable[[T, int], str]:
    """
    Return a json serializer that serializes the json msg, adds the
    associated version, and then serializes the RedisState proto to a string
    """
    def _serialize_json(msg: T, version: int) -> str:
        serialized_msg = jsonpickle.encode(msg)
        redis_state = RedisState(
            serialized_msg=serialized_msg.encode('utf-8'),
            version=version,
            is_garbage=False,
        )
        return redis_state.SerializeToString()

    return _serialize_json


def get_json_deserializer() -> Callable[[str], T]:
    """
    Returns a json deserializer that deserializes the RedisState proto and
    then deserializes the json msg
    """
    def _deserialize_json(serialized_rule: str) -> T:
        proto_wrapper = RedisState()
        proto_wrapper.ParseFromString(serialized_rule)
        serialized_msg = proto_wrapper.serialized_msg
        msg = jsonpickle.decode(serialized_msg.decode('utf-8'))
        return msg

    return _deserialize_json


def get_proto_version_deserializer() -> Callable[[str], T]:
    """
    Return a proto deserializer that takes in a proto type to deserialize
    the version number stored in the RedisState proto
    """
    def _deserialize_version(serialized_rule: str) -> T:
        proto_wrapper = RedisState()
        proto_wrapper.ParseFromString(serialized_rule)
        return proto_wrapper.version
    return _deserialize_version