privacysandbox / protected-auction-key-value-service

Protected Auction Key/Value Service
Apache License 2.0
55 stars 20 forks source link

Sample for writing data loading DataRecord using Java #39

Closed kelvintatendagorekore closed 8 months ago

kelvintatendagorekore commented 11 months ago

Please note that String type in our flatbuffer schema has been renamed to StringType in the following code because the current name conflicts with Java native type java.lang.String and we are working on fixing this.

package kv_server;

import com.google.flatbuffers.FlatBufferBuilder;
import java.nio.ByteBuffer;

public class DataRecordExample {
    static final String KEY = "record key";
    static final String VALUE = "record value";
    static final long LOGICAL_COMMIT_TIME = 123456789;

    static int createStringType(FlatBufferBuilder builder, String value) {
        int valueString = builder.createString(value);
        return StringType.createStringType(builder, valueString);
    }

    static byte[] createKeyValueMutationDataRecord() {
        FlatBufferBuilder builder = new FlatBufferBuilder();
        int key = builder.createString(KEY);
        int stringType = createStringType(builder, VALUE);
        KeyValueMutationRecord.startKeyValueMutationRecord(builder);
        KeyValueMutationRecord.addKey(builder, key);
        KeyValueMutationRecord.addValueType(builder, Value.StringType);
        KeyValueMutationRecord.addValue(builder, stringType);
        KeyValueMutationRecord.addMutationType(builder, KeyValueMutationType.Update);
        KeyValueMutationRecord.addLogicalCommitTime(builder, LOGICAL_COMMIT_TIME);
        int mutationRecord = KeyValueMutationRecord.endKeyValueMutationRecord(builder);

        DataRecord.startDataRecord(builder);
        DataRecord.addRecord(builder, mutationRecord);
        DataRecord.addRecordType(builder, Record.KeyValueMutationRecord);
        int dataRecord = DataRecord.endDataRecord(builder);
        builder.finish(dataRecord); // We are done writing so finish the buffer.
        // Serialize to byte buffer and return
        return builder.sizedByteArray();
    }

    static DataRecord readKeyValueMutationDataRecord(byte[] buffer) {
        return DataRecord.getRootAsDataRecord(ByteBuffer.wrap(buffer));
    }

    public static void main(String[] args) {
        byte[] dataRecordBuffer = createKeyValueMutationDataRecord();
        DataRecord dataRecord = readKeyValueMutationDataRecord(dataRecordBuffer);
        int recordType = dataRecord.recordType();
        System.out.println("Record type: " + dataRecord.recordType());
        if (recordType != Record.KeyValueMutationRecord) {
            throw new RuntimeException("Failed to serialize and deserialize data record.");
        }
        KeyValueMutationRecord keyValueMutationRecord = (KeyValueMutationRecord) dataRecord
                .record(new KeyValueMutationRecord());
        System.out.println("Record key: " + keyValueMutationRecord.key());
        System.out.println("Record timestamp: " + keyValueMutationRecord.logicalCommitTime());
        System.out.println("Record mutation type: " + keyValueMutationRecord.mutationType());
        if (keyValueMutationRecord.valueType() != Value.StringType) {
            throw new RuntimeException("Failed to serialize and deserialize key value mutation record.");
        }
        StringType value = (StringType) keyValueMutationRecord.value(new StringType());
        System.out.println("Record value: " + value.value());        
    }
}