From 54480806afb3b7f0886a7839146045537f9d95d9 Mon Sep 17 00:00:00 2001 From: Jonathan Flynn Date: Wed, 25 Sep 2024 15:32:56 +0100 Subject: [PATCH 1/5] add methods for reading and writing enums --- .../com/twitter/scrooge/LazyTProtocol.scala | 2 ++ .../twitter/scrooge/TLazyBinaryProtocol.scala | 17 ++++++++++++++++- .../com/twitter/scrooge/backend/Generator.scala | 1 + 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala b/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala index 3951150b0..240b1e7c3 100644 --- a/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala +++ b/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala @@ -124,4 +124,6 @@ trait LazyTProtocol extends TProtocol { */ def offsetSkipBinary(): Int + def writeEnum(value: Int): Unit + def readEnum(): Int } diff --git a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala index 33968fcd3..7938cb0c2 100644 --- a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala +++ b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala @@ -141,6 +141,11 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) } } + override def writeEnum(value: Int): Unit = { + writeByte(-1.toByte) + writeI32(value) + } + override def writeBinary(bin: ByteBuffer): Unit = { val length: Int = bin.limit() - bin.position() - bin.arrayOffset() val buf = transport.getBuffer(length + 4) @@ -176,7 +181,8 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) override def readFieldBegin(): TField = { val tpe: Byte = readByte() val id: Short = if (tpe == TType.STOP) 0 else readI16() - new TField("", tpe, id) + val finalType = if (tpe == -1) TType.ENUM else tpe + new TField("", finalType, id) } override def readFieldEnd(): Unit = () @@ -255,6 +261,15 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) throw new TException("JVM DOES NOT SUPPORT UTF-8") } + override def readEnum(): Int = { + val typeId = readByte() + if (typeId == -1 || typeId == 16) { // Handle both old and new ENUM identifiers, see PR link for more information. + readI32() + } else { + throw new TException(s"Invalid type for enum: $typeId") + } + } + override def buffer: Array[Byte] = transport.srcBuf override def offset: Int = transport.getBufferPosition diff --git a/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala b/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala index 84e5eb4f1..269e53115 100644 --- a/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala +++ b/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala @@ -489,6 +489,7 @@ abstract class TemplateGenerator(val resolvedDoc: ResolvedDocument) case TDouble => "readDouble" case TString => "readString" case TBinary => "readBinary" + case EnumType(_,_) => "readI32" case x => throw new ScroogeInternalException("genProtocolReadMethod#" + t) } v(getCode(t)) From 8c4f7f28ab2c4d3c817277da6a14d151b6865c14 Mon Sep 17 00:00:00 2001 From: Jonathan Flynn Date: Wed, 25 Sep 2024 15:44:14 +0100 Subject: [PATCH 2/5] create variables for enum types --- .../main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala | 4 +++- .../main/scala/com/twitter/scrooge/backend/Generator.scala | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala index 7938cb0c2..e053b1e4e 100644 --- a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala +++ b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala @@ -17,6 +17,8 @@ import org.apache.thrift.protocol._ object TLazyBinaryProtocol { private val AnonymousStruct: TStruct = new TStruct() private val utf8Charset = Charset.forName("UTF-8") + private val OLD_ENUM_TYPE_ID: Byte = 16 + private val NEW_ENUM_TYPE_ID: Byte = -1 } class TLazyBinaryProtocol(transport: TArrayByteTransport) @@ -263,7 +265,7 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) override def readEnum(): Int = { val typeId = readByte() - if (typeId == -1 || typeId == 16) { // Handle both old and new ENUM identifiers, see PR link for more information. + if (typeId == NEW_ENUM_TYPE_ID || typeId == OLD_ENUM_TYPE_ID) { // Handle both old and new ENUM identifiers, see PR link for more information. readI32() } else { throw new TException(s"Invalid type for enum: $typeId") diff --git a/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala b/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala index 269e53115..80be8ed2b 100644 --- a/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala +++ b/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala @@ -560,6 +560,7 @@ abstract class TemplateGenerator(val resolvedDoc: ResolvedDocument) case TDouble => "writeDouble" case TString => "writeString" case TBinary => "writeBinary" + case EnumType(_,_) => "writeI32" case x => throw new ScroogeInternalException("protocolWriteMethod#" + t) } v(getCode(t)) From 6f30455e2ebf566e6559417d45b86580106cea86 Mon Sep 17 00:00:00 2001 From: Jonathan Flynn Date: Wed, 25 Sep 2024 15:45:27 +0100 Subject: [PATCH 3/5] create variables for enum types --- .../main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala index e053b1e4e..b78a47e05 100644 --- a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala +++ b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala @@ -183,7 +183,7 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) override def readFieldBegin(): TField = { val tpe: Byte = readByte() val id: Short = if (tpe == TType.STOP) 0 else readI16() - val finalType = if (tpe == -1) TType.ENUM else tpe + val finalType = if (tpe == NEW_ENUM_TYPE_ID) TType.ENUM else tpe new TField("", finalType, id) } From 125f3ea1662e36778be5755e42e99914538ff7f6 Mon Sep 17 00:00:00 2001 From: Jonathan Flynn Date: Wed, 25 Sep 2024 16:01:46 +0100 Subject: [PATCH 4/5] remove redundant methods --- .../com/twitter/scrooge/LazyTProtocol.scala | 3 --- .../twitter/scrooge/TLazyBinaryProtocol.scala | 18 ++---------------- 2 files changed, 2 insertions(+), 19 deletions(-) diff --git a/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala b/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala index 240b1e7c3..e261cfbb4 100644 --- a/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala +++ b/scrooge-core/src/main/scala/com/twitter/scrooge/LazyTProtocol.scala @@ -123,7 +123,4 @@ trait LazyTProtocol extends TProtocol { * Returns: The offset at which the string can be read. */ def offsetSkipBinary(): Int - - def writeEnum(value: Int): Unit - def readEnum(): Int } diff --git a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala index b78a47e05..3dbcc37a4 100644 --- a/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala +++ b/scrooge-core/src/main/scala/com/twitter/scrooge/TLazyBinaryProtocol.scala @@ -17,7 +17,6 @@ import org.apache.thrift.protocol._ object TLazyBinaryProtocol { private val AnonymousStruct: TStruct = new TStruct() private val utf8Charset = Charset.forName("UTF-8") - private val OLD_ENUM_TYPE_ID: Byte = 16 private val NEW_ENUM_TYPE_ID: Byte = -1 } @@ -32,9 +31,10 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) } override def writeFieldBegin(field: TField): Unit = { + val typeToWrite = if (field.`type` == TType.ENUM) NEW_ENUM_TYPE_ID else field.`type` val buf = transport.getBuffer(3) val offset = transport.writerOffset - buf(offset) = field.`type` + buf(offset) = typeToWrite innerWriteI16(buf, offset + 1, field.id) } @@ -143,11 +143,6 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) } } - override def writeEnum(value: Int): Unit = { - writeByte(-1.toByte) - writeI32(value) - } - override def writeBinary(bin: ByteBuffer): Unit = { val length: Int = bin.limit() - bin.position() - bin.arrayOffset() val buf = transport.getBuffer(length + 4) @@ -263,15 +258,6 @@ class TLazyBinaryProtocol(transport: TArrayByteTransport) throw new TException("JVM DOES NOT SUPPORT UTF-8") } - override def readEnum(): Int = { - val typeId = readByte() - if (typeId == NEW_ENUM_TYPE_ID || typeId == OLD_ENUM_TYPE_ID) { // Handle both old and new ENUM identifiers, see PR link for more information. - readI32() - } else { - throw new TException(s"Invalid type for enum: $typeId") - } - } - override def buffer: Array[Byte] = transport.srcBuf override def offset: Int = transport.getBufferPosition From 4d7459fa1ab4d3132937ae33353187816514689a Mon Sep 17 00:00:00 2001 From: Jonathan Flynn Date: Thu, 26 Sep 2024 09:49:45 +0100 Subject: [PATCH 5/5] add unit tests for reading old and new enums --- .../scrooge/internal/TProtocolsTest.scala | 45 +++++++++++++++++-- .../twitter/scrooge/backend/Generator.scala | 2 - 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/scrooge-core/src/test/scala/com/twitter/scrooge/internal/TProtocolsTest.scala b/scrooge-core/src/test/scala/com/twitter/scrooge/internal/TProtocolsTest.scala index f33803dbd..7ab733ed9 100644 --- a/scrooge-core/src/test/scala/com/twitter/scrooge/internal/TProtocolsTest.scala +++ b/scrooge-core/src/test/scala/com/twitter/scrooge/internal/TProtocolsTest.scala @@ -1,8 +1,6 @@ package com.twitter.scrooge.internal -import com.twitter.scrooge.TArrayByteTransport -import com.twitter.scrooge.TFieldBlob -import com.twitter.scrooge.ThriftUnion +import com.twitter.scrooge.{TArrayByteTransport, TFieldBlob, TLazyBinaryProtocol, ThriftUnion} import com.twitter.util.mock.Mockito import org.apache.thrift.protocol.TBinaryProtocol import org.apache.thrift.protocol.TCompactProtocol @@ -10,6 +8,7 @@ import org.apache.thrift.protocol.TField import org.apache.thrift.protocol.TProtocolException import org.apache.thrift.protocol.TType import org.apache.thrift.transport.TMemoryBuffer + import scala.collection.immutable import org.scalatest.funsuite.AnyFunSuite @@ -173,4 +172,44 @@ class TProtocolsTest extends AnyFunSuite with Mockito { succeed } + test("readFieldBegin handles new ENUM type identifier") { + val writeBuffer = new TMemoryBuffer(128) + val writeProto = new TBinaryProtocol(writeBuffer) + + val newEnum = -1 + + writeProto.writeByte(newEnum.toByte) // New ENUM type identifier + writeProto.writeI16(1) // Field ID + writeProto.writeI32(2) // Enum value + + val readBuffer = TArrayByteTransport(writeBuffer.getArray) + val readProto = new TLazyBinaryProtocol(readBuffer) + + val field = readProto.readFieldBegin() + assert(field.`type` == TType.ENUM) + assert(field.id == 1) + + val enumValue = readProto.readI32() + assert(enumValue == 2) + } + + test("readFieldBegin handles old ENUM type identifier") { + val writeBuffer = new TMemoryBuffer(128) + val writeProto = new TBinaryProtocol(writeBuffer) + + writeProto.writeByte(TType.ENUM) + writeProto.writeI16(1) + writeProto.writeI32(2) + + val readBuffer = TArrayByteTransport(writeBuffer.getArray) + val readProto = new TLazyBinaryProtocol(readBuffer) + + val field = readProto.readFieldBegin() + assert(field.`type` == TType.ENUM) + assert(field.id == 1) + + val enumValue = readProto.readI32() + assert(enumValue == 2) + } + } diff --git a/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala b/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala index 80be8ed2b..84e5eb4f1 100644 --- a/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala +++ b/scrooge-generator/src/main/scala/com/twitter/scrooge/backend/Generator.scala @@ -489,7 +489,6 @@ abstract class TemplateGenerator(val resolvedDoc: ResolvedDocument) case TDouble => "readDouble" case TString => "readString" case TBinary => "readBinary" - case EnumType(_,_) => "readI32" case x => throw new ScroogeInternalException("genProtocolReadMethod#" + t) } v(getCode(t)) @@ -560,7 +559,6 @@ abstract class TemplateGenerator(val resolvedDoc: ResolvedDocument) case TDouble => "writeDouble" case TString => "writeString" case TBinary => "writeBinary" - case EnumType(_,_) => "writeI32" case x => throw new ScroogeInternalException("protocolWriteMethod#" + t) } v(getCode(t))