[ https://issues.apache.org/jira/browse/BEAM-7274?focusedWorklogId=362375&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-362375 ]
ASF GitHub Bot logged work on BEAM-7274: ---------------------------------------- Author: ASF GitHub Bot Created on: 23/Dec/19 01:12 Start Date: 23/Dec/19 01:12 Worklog Time Spent: 10m Work Description: reuvenlax commented on pull request #10449: [BEAM-7274] Implement the Protobuf schema provider for compiled protocol buffers URL: https://github.com/apache/beam/pull/10449#discussion_r360741862 ########## File path: sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoByteBuddyUtils.java ########## @@ -0,0 +1,573 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.extensions.protobuf; + +import static org.apache.beam.sdk.extensions.protobuf.ProtoSchemaTranslator.getFieldNumber; + +import com.google.protobuf.ByteString; +import com.google.protobuf.Duration; +import com.google.protobuf.Internal.EnumLite; +import com.google.protobuf.MessageLite; +import com.google.protobuf.ProtocolMessageEnum; +import com.google.protobuf.Timestamp; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Type; +import java.util.List; +import java.util.Map; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.beam.sdk.extensions.protobuf.ProtoSchemaLogicalTypes.DurationNanos; +import org.apache.beam.sdk.extensions.protobuf.ProtoSchemaLogicalTypes.TimestampNanos; +import org.apache.beam.sdk.schemas.FieldValueGetter; +import org.apache.beam.sdk.schemas.FieldValueSetter; +import org.apache.beam.sdk.schemas.FieldValueTypeInformation; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.Field; +import org.apache.beam.sdk.schemas.Schema.FieldType; +import org.apache.beam.sdk.schemas.Schema.TypeName; +import org.apache.beam.sdk.schemas.SchemaUserTypeCreator; +import org.apache.beam.sdk.schemas.logicaltypes.OneOfType; +import org.apache.beam.sdk.schemas.logicaltypes.OneOfType.Value; +import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConvertType; +import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConvertValueForGetter; +import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConvertValueForSetter; +import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.InjectPackageStrategy; +import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.TypeConversion; +import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.TypeConversionsFactory; +import org.apache.beam.sdk.schemas.utils.FieldValueTypeSupplier; +import org.apache.beam.sdk.schemas.utils.JavaBeanUtils; +import org.apache.beam.sdk.schemas.utils.ReflectUtils; +import org.apache.beam.sdk.schemas.utils.ReflectUtils.ClassWithSchema; +import org.apache.beam.sdk.util.common.ReflectHelpers; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.ByteBuddy; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.description.type.TypeDescription.ForLoadedType; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.dynamic.DynamicType; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.dynamic.loading.ClassLoadingStrategy; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.dynamic.scaffold.InstrumentedType; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.Implementation; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.ByteCodeAppender; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.StackManipulation; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.StackManipulation.Compound; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.assign.Assigner; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.assign.Assigner.Typing; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.assign.TypeCasting; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.member.MethodInvocation; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.implementation.bytecode.member.MethodReturn; +import org.apache.beam.vendor.bytebuddy.v1_9_3.net.bytebuddy.matcher.ElementMatchers; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.CaseFormat; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Multimap; + +public class ProtoByteBuddyUtils { + private static final ByteBuddy BYTE_BUDDY = new ByteBuddy(); + private static TypeDescriptor<ByteString> BYTE_STRING_TYPE_DESCRIPTOR = + TypeDescriptor.of(ByteString.class); + private static TypeDescriptor<Timestamp> PROTO_TIMESTAMP_TYPE_DESCRIPTOR = + TypeDescriptor.of(Timestamp.class); + private static TypeDescriptor<Duration> PROTO_DURATION_TYPE_DESCRIPTOR = + TypeDescriptor.of(Duration.class); + + private static final ForLoadedType BYTE_STRING_TYPE = new ForLoadedType(ByteString.class); + private static final ForLoadedType BYTE_ARRAY_TYPE = new ForLoadedType(byte[].class); + private static final ForLoadedType PROTO_ENUM_TYPE = new ForLoadedType(ProtocolMessageEnum.class); + private static final ForLoadedType INTEGER_TYPE = new ForLoadedType(Integer.class); + + private static final Map<TypeName, String> PROTO_GETTER_SUFFIX = + ImmutableMap.of( + TypeName.ARRAY, "List", + TypeName.ITERABLE, "List", + TypeName.MAP, "Map"); + + private static final Map<TypeName, String> PROTO_SETTER_PREFIX = + ImmutableMap.of( + TypeName.ARRAY, "addAll", + TypeName.ITERABLE, "addAll", + TypeName.MAP, "putAll"); + private static final String DEFAULT_PROTO_GETTER_PREFIX = "get"; + private static final String DEFAULT_PROTO_SETTER_PREFIX = "set"; + + static String protoGetterName(String name, FieldType fieldType) { + final String camel = CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, name); + return DEFAULT_PROTO_GETTER_PREFIX + + camel + + PROTO_GETTER_SUFFIX.getOrDefault(fieldType.getTypeName(), ""); + } + + static String protoSetterName(String name, FieldType fieldType) { + final String camel = CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, name); + return protoSetterPrefix(fieldType) + camel; + } + + static String protoSetterPrefix(FieldType fieldType) { + return PROTO_SETTER_PREFIX.getOrDefault(fieldType.getTypeName(), DEFAULT_PROTO_SETTER_PREFIX); + } + + static class ProtoConvertType extends ConvertType { + ProtoConvertType(boolean returnRawValues) { + super(returnRawValues); + } + + @Override + public Type convert(TypeDescriptor typeDescriptor) { + if (typeDescriptor.equals(BYTE_STRING_TYPE_DESCRIPTOR) + || typeDescriptor.isSubtypeOf(BYTE_STRING_TYPE_DESCRIPTOR)) { + return byte[].class; + } else if (typeDescriptor.isSubtypeOf(TypeDescriptor.of(ProtocolMessageEnum.class))) { + return Integer.class; + } else if (typeDescriptor.equals(PROTO_TIMESTAMP_TYPE_DESCRIPTOR) + || typeDescriptor.equals(PROTO_DURATION_TYPE_DESCRIPTOR)) { + return Row.class; + } else { + return super.convert(typeDescriptor); + } + } + } + + static class ProtoConvertValueForGetter extends ConvertValueForGetter { + ProtoConvertValueForGetter(StackManipulation readValue) { + super(readValue); + } + + @Override + protected ProtoTypeConversionsFactory getFactory() { + return new ProtoTypeConversionsFactory(); + } + + @Override + public StackManipulation convert(TypeDescriptor type) { + if (type.equals(BYTE_STRING_TYPE_DESCRIPTOR) + || type.isSubtypeOf(BYTE_STRING_TYPE_DESCRIPTOR)) { + return new Compound( + readValue, + MethodInvocation.invoke( + BYTE_STRING_TYPE + .getDeclaredMethods() + .filter(ElementMatchers.named("toByteArray")) + .getOnly())); + } else if (type.isSubtypeOf(TypeDescriptor.of(ProtocolMessageEnum.class))) { + return new Compound( + readValue, + MethodInvocation.invoke( + PROTO_ENUM_TYPE + .getDeclaredMethods() + .filter( + ElementMatchers.named("getNumber").and(ElementMatchers.takesArguments(0))) + .getOnly()), + Assigner.DEFAULT.assign( + INTEGER_TYPE.asUnboxed().asGenericType(), + INTEGER_TYPE.asGenericType(), + Typing.STATIC)); + } else if (type.equals(PROTO_TIMESTAMP_TYPE_DESCRIPTOR)) { + return new Compound( + readValue, + MethodInvocation.invoke( + new ForLoadedType(TimestampNanos.class) + .getDeclaredMethods() + .filter(ElementMatchers.named("toRow")) + .getOnly())); + } else if (type.equals(PROTO_DURATION_TYPE_DESCRIPTOR)) { + return new Compound( + readValue, + MethodInvocation.invoke( + new ForLoadedType(DurationNanos.class) + .getDeclaredMethods() + .filter(ElementMatchers.named("toRow")) + .getOnly())); + } else { + return super.convert(type); + } + } + } + + static class ProtoConvertValueForSetter extends ConvertValueForSetter { + ProtoConvertValueForSetter(StackManipulation readValue) { + super(readValue); + } + + @Override + protected ProtoTypeConversionsFactory getFactory() { + return new ProtoTypeConversionsFactory(); + } + + @Override + public StackManipulation convert(TypeDescriptor type) { + if (type.isSubtypeOf(TypeDescriptor.of(ByteString.class))) { + return new Compound( + readValue, + TypeCasting.to(BYTE_ARRAY_TYPE), + MethodInvocation.invoke( + BYTE_STRING_TYPE + .getDeclaredMethods() + .filter( + ElementMatchers.named("copyFrom") + .and(ElementMatchers.takesArguments(BYTE_ARRAY_TYPE))) + .getOnly())); + } else if (type.isSubtypeOf(TypeDescriptor.of(ProtocolMessageEnum.class))) { + ForLoadedType loadedType = new ForLoadedType(type.getRawType()); + // Convert the stored number back to the enum constant. + return new Compound( + readValue, + Assigner.DEFAULT.assign( + INTEGER_TYPE.asBoxed().asGenericType(), + INTEGER_TYPE.asUnboxed().asGenericType(), + Typing.STATIC), + MethodInvocation.invoke( + loadedType + .getDeclaredMethods() + .filter( + ElementMatchers.named("forNumber") + .and(ElementMatchers.isStatic().and(ElementMatchers.takesArguments(1)))) + .getOnly())); + } else if (type.equals(PROTO_TIMESTAMP_TYPE_DESCRIPTOR)) { + return new Compound( + readValue, + MethodInvocation.invoke( + new ForLoadedType(TimestampNanos.class) + .getDeclaredMethods() + .filter(ElementMatchers.named("toDuration")) + .getOnly())); + } else if (type.equals(PROTO_DURATION_TYPE_DESCRIPTOR)) { + return new Compound( + readValue, + MethodInvocation.invoke( + new ForLoadedType(DurationNanos.class) + .getDeclaredMethods() + .filter(ElementMatchers.named("toDuration")) + .getOnly())); + } else { Review comment: yes! that's why tests are failing, and I'll add that support. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org Issue Time Tracking ------------------- Worklog Id: (was: 362375) Time Spent: 21h (was: 20h 50m) > Protobuf Beam Schema support > ---------------------------- > > Key: BEAM-7274 > URL: https://issues.apache.org/jira/browse/BEAM-7274 > Project: Beam > Issue Type: Improvement > Components: sdk-java-core > Reporter: Alex Van Boxel > Assignee: Alex Van Boxel > Priority: Minor > Time Spent: 21h > Remaining Estimate: 0h > > Add support for the new Beam Schema to the Protobuf extension. -- This message was sent by Atlassian Jira (v8.3.4#803005)