Attach bonus properties to a protobuf message

person.proto

// From https://github.com/square/wire/blob/master/wire-tests/src/commonTest/proto/kotlin/person.proto

/*
 * Copyright 2013 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package squareup.protos.kotlin.person;

option java_package = "com.squareup.wire.protos.kotlin.person";
option java_outer_classname = "PersonProtos";

/**
 * Message representing a person, includes their name, unique ID number, email and phone number.
 */
message Person {
  // The customer's ID number.
  required int32 id = 2;
  // The customer's full name.
  required string name = 1;
  // Email address for the customer.
  optional string email = 3;

  /**
   * Represents the type of the phone number: mobile, home or work.
   */
  enum PhoneType {
    MOBILE = 0;
    HOME = 1;
    /** Could be phone or fax. */
    WORK = 2;
  }

  message PhoneNumber {
    // The customer's phone number.
    required string number = 1;
    // The type of phone stored here.
    optional PhoneType type = 2 [default = HOME];
  }

  // A list of the customer's phone numbers.
  repeated PhoneNumber phone = 4;

  repeated string aliases = 5;
}

RuntimePropertiesDecoder.kt

package com.squareup.wire.samples

import com.squareup.wire.Message
import com.squareup.wire.ProtoAdapter
import com.squareup.wire.ProtoReader32
import com.squareup.wire.ProtoWriter
import com.squareup.wire.forEachTag
import okio.Buffer

/**
 * Read and write properties to a message type by tag, even if those
 * properties aren't declared in the schema.
 *
 * This requires that types are unique among tags.
 */
class RuntimePropertiesDecoder<M : Message<M, *>>(
  private val tagToAdapter: Map<Int, ProtoAdapter<*>>,
) {
  private val classToEntry = tagToAdapter.entries.associateBy { it.value.type }

  init {
    require(classToEntry.size == tagToAdapter.size) {
      "adapter types must be unique"
    }
  }

  fun get(message: M): Any? {
    val reader = ProtoReader32(message.unknownFields)
    reader.forEachTag { tag ->
      val adapter = tagToAdapter[tag]
      when {
        adapter != null -> return@get adapter.decode(reader)
        else -> reader.readUnknownField(tag)
      }
    }
    return null
  }

  fun put(message: M, value: Any): M? {
    val targetClass = value::class
    val (tag, adapter) = classToEntry[targetClass] ?: return null // No registered type

    val buffer = Buffer()
    val writer = ProtoWriter(buffer)
    message.adapter.encode(writer, message)
    (adapter as ProtoAdapter<Any>).encodeWithTag(writer, tag, value)
    return message.adapter.decode(buffer)
  }
}

RuntimePropertiesTest.kt

package com.squareup.wire.samples

import assertk.assertThat
import assertk.assertions.isEqualTo
import com.squareup.wire.ProtoAdapter
import com.squareup.wire.protos.kotlin.person.Person
import kotlin.test.Test

class RuntimePropertiesTest {
  @Test
  fun happyPath() {
    val runtimePropertiesDecoder = RuntimePropertiesDecoder<Person>(
      mapOf(
        100 to ProtoAdapter.STRING,
        101 to ProtoAdapter.BYTES,
        102 to Person.PhoneNumber.ADAPTER,
      ),
    )

    val person = Person(
      id = 1,
      name = "Jesse",
    )

    val personWithCity = runtimePropertiesDecoder.put(person, "Kitchener")!!
    assertThat(runtimePropertiesDecoder.get(personWithCity)).isEqualTo("Kitchener")

    val phoneNumber = Person.PhoneNumber(
      number = "519-867-5309",
      type = Person.PhoneType.MOBILE,
    )
    val personWithPhoneNumber = runtimePropertiesDecoder.put(
      person,
      phoneNumber,
    )!!
    assertThat(runtimePropertiesDecoder.get(personWithPhoneNumber)).isEqualTo(phoneNumber)
  }
}

Usually unknownFields tracks data that the current schema doesn’t have yet, but it’s perfectly valid to use unknown fields to track data that the current schema doesn’t have any more.

Use this to remove imports from your .proto declarations, without losing the ability to read & write those messages at runtime.