LinyuYao1021 commented on a change in pull request #14737: URL: https://github.com/apache/flink/pull/14737#discussion_r573954198
########## File path: flink-formats/flink-avro-glue-schema-registry/src/main/java/org/apache/flink/formats/avro/glue/schema/registry/GlueSchemaRegistryAvroSerializationSchema.java ########## @@ -0,0 +1,127 @@ +/* + * 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.flink.formats.avro.glue.schema.registry; + +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.formats.avro.RegistryAvroSerializationSchema; +import org.apache.flink.formats.avro.SchemaCoder; + +import lombok.SneakyThrows; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.Encoder; +import org.apache.avro.specific.SpecificRecord; + +import javax.annotation.Nullable; + +import java.io.ByteArrayOutputStream; +import java.util.Map; + +/** + * AWS Glue Schema Registry Serialization schema to serialize to Avro binary format for Flink + * Producer user. + * + * @param <T> the type to be serialized + */ +public class GlueSchemaRegistryAvroSerializationSchema<T> + extends RegistryAvroSerializationSchema<T> { + /** + * Creates an Avro serialization schema. + * + * @param recordClazz class to serialize. Should be one of: {@link SpecificRecord}, {@link + * GenericRecord}. + * @param reader reader's Avro schema. Should be provided if recordClazz is {@link + * GenericRecord} + * @param schemaCoderProvider schema coder provider which reads writer schema from AWS Glue + * Schema Registry + */ + private GlueSchemaRegistryAvroSerializationSchema( + Class<T> recordClazz, + @Nullable Schema reader, + SchemaCoder.SchemaCoderProvider schemaCoderProvider) { + super(recordClazz, reader, schemaCoderProvider); + } + + @VisibleForTesting + protected GlueSchemaRegistryAvroSerializationSchema( + Class<T> recordClazz, @Nullable Schema reader, SchemaCoder schemaCoder) { + // Pass null schema coder provider + super(recordClazz, reader, null); + this.schemaCoder = schemaCoder; + } + + /** + * Creates {@link GlueSchemaRegistryAvroSerializationSchema} that serializes {@link + * GenericRecord} using provided schema. + * + * @param schema the schema that will be used for serialization + * @param transportName topic name or stream name etc. + * @param configs configuration map of AWS Glue Schema Registry + * @return serialized record in form of byte array + */ + public static GlueSchemaRegistryAvroSerializationSchema<GenericRecord> forGeneric( + Schema schema, String transportName, Map<String, Object> configs) { + return new GlueSchemaRegistryAvroSerializationSchema<>( + GenericRecord.class, + schema, + new GlueSchemaRegistryAvroSchemaCoderProvider(transportName, configs)); + } + + /** + * Creates {@link GlueSchemaRegistryAvroSerializationSchema} that serializes {@link + * SpecificRecord} using provided schema. + * + * @param clazz the type to be serialized + * @param transportName topic name or stream name etc. + * @param configs configuration map of Amazon Schema Registry + * @return serialized record in form of byte array + */ + public static <T extends SpecificRecord> + GlueSchemaRegistryAvroSerializationSchema<T> forSpecific( + Class<T> clazz, String transportName, Map<String, Object> configs) { + return new GlueSchemaRegistryAvroSerializationSchema<>( + clazz, null, new GlueSchemaRegistryAvroSchemaCoderProvider(transportName, configs)); + } + + /** + * Serializes the incoming element to a byte array containing bytes of AWS Glue Schema registry + * information. + * + * @param object The incoming element to be serialized + * @return The serialized bytes. + */ + @SneakyThrows + @Override + public byte[] serialize(T object) { + checkAvroInitialized(); + + if (object == null) { + return null; + } else { + ByteArrayOutputStream outputStream = getOutputStream(); + outputStream.reset(); + Encoder encoder = getEncoder(); + getDatumWriter().write(object, encoder); + schemaCoder.writeSchema(getSchema(), outputStream); + encoder.flush(); + + return outputStream.toByteArray(); + } + } Review comment: Yes, the reason to override this method is that it's required to add header byte and compression byte in the serialized byte array. So we change the order of writing schema and writing object. And we have extra methods in _writeSchema_ to achieve this feature. ---------------------------------------------------------------- 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
