如何配置PyMongo以自动编解码Protobuf消息

需求概述 #

• 允许在PyMongo所有支持传入dict的场景（如collection.insert_one()）直接传入Protobuf消息实例
• 在返回dict的场景（如collection.find_one()）直接获取对应的Protobuf实例

现有方案的痛点 #

当前通过json_format.MessageToDict()做双向转换的方式存在以下问题：

• 双重转换（Protobuf→dict→BSON / BSON→dict→Protobuf）效率低下
• json_format生成JSON安全的dict时会丢失类型信息：bytes被转成base64、int64被转成字符串，导致数据失真
• 二进制等特殊类型处理存在兼容性问题

可行解决方案：自定义BSON编解码器 #

PyMongo的TypeCodec虽多用于单个类型，但可通过自定义CodecRegistry实现对整个Protobuf消息类型的编解码，核心是直接在Protobuf与BSON间转换，跳过dict中间层。

步骤1：实现Protobuf编解码器 #

编写继承TypeCodec的类，针对Protobuf的Message类型实现编码、解码逻辑：

from bson import TypeCodec, TypeRegistry, CodecRegistry
from google.protobuf.message import Message

class ProtobufCodec(TypeCodec):
    def __init__(self, message_type: type[Message]):
        self.message_type = message_type
        self.encoder_type = Message
        self.decoder_type = dict

    def transform_python(self, value: Message) -> dict:
        """将Protobuf消息转为BSON可序列化结构"""
        return {
            "_pb_type": self.message_type.DESCRIPTOR.full_name,
            "_pb_data": value.SerializeToString()
        }

    def transform_bson(self, value: dict) -> Message:
        """将BSON结构转为Protobuf消息实例"""
        if "_pb_data" in value and "_pb_type" in value:
            msg = self.message_type()
            msg.ParseFromString(value["_pb_data"])
            return msg
        return value

步骤2：注册编解码器到PyMongo客户端 #

将编解码器绑定到MongoClient或Collection：

from pymongo import MongoClient
from your_proto_module import YourProtobufMessage  # 替换为你的Protobuf类型

# 注册编解码器
pb_codec = ProtobufCodec(YourProtobufMessage)
type_registry = TypeRegistry([pb_codec])
codec_registry = CodecRegistry(type_registry=type_registry)

# 创建带编解码器的客户端
client = MongoClient(codec_registry=codec_registry)
db = client["your_db"]
collection = db["your_collection"]

步骤3：直接使用Protobuf实例操作 #

# 插入Protobuf消息
msg = YourProtobufMessage(field1="value1", field2=123)
collection.insert_one(msg)

# 查询并直接获取Protobuf实例
result = collection.find_one({"field1": "value1"})
print(isinstance(result, YourProtobufMessage))  # 输出True

进阶：支持多Protobuf类型 #

若需处理多种消息类型，可在编解码器中维护类型映射，通过_pb_type字段动态匹配：

class MultiProtobufCodec(TypeCodec):
    def __init__(self, type_map: dict[str, type[Message]]):
        self.type_map = type_map
        self.encoder_type = Message
        self.decoder_type = dict

    def transform_python(self, value: Message) -> dict:
        return {
            "_pb_type": value.DESCRIPTOR.full_name,
            "_pb_data": value.SerializeToString()
        }

    def transform_bson(self, value: dict) -> Message:
        if "_pb_data" in value and "_pb_type" in value:
            msg_type = self.type_map.get(value["_pb_type"])
            if msg_type:
                msg = msg_type()
                msg.ParseFromString(value["_pb_data"])
                return msg
        return value

# 使用示例
type_map = {
    "your_proto.YourProtobufMessage": YourProtobufMessage,
    "your_proto.AnotherMessage": AnotherMessage
}
multi_codec = MultiProtobufCodec(type_map)
codec_registry = CodecRegistry(type_registry=TypeRegistry([multi_codec]))

注意事项 #

• 存储的BSON文档会包含_pb_type（标识消息类型）和_pb_data（存储二进制数据）两个字段
• 如需兼容原有dict格式存储的数据，可在解码逻辑中增加判断：无_pb_type字段时返回原始dict
• 该方式直接操作Protobuf二进制序列化，避免了中间转换的性能损耗和类型失真问题

内容的提问来源于stack exchange，提问作者James S