forked from apache/pulsar
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Kinesis sink publish full json message (apache#2079)
* Kinesis sink publish full json message * fix pulsar typo
- Loading branch information
Showing
5 changed files
with
324 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
103 changes: 103 additions & 0 deletions
103
pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/Utils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,103 @@ | ||
/** | ||
* 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.pulsar.io.kinesis; | ||
|
||
import static java.util.Base64.getEncoder; | ||
|
||
import java.util.Map; | ||
|
||
import org.apache.pulsar.common.api.EncryptionContext; | ||
import org.apache.pulsar.io.core.RecordContext; | ||
|
||
import com.google.gson.JsonObject; | ||
|
||
public class Utils { | ||
|
||
private static final String PAYLOAD_FIELD = "payloadBase64"; | ||
private static final String PROPERTIES_FIELD = "properties"; | ||
private static final String KEY_MAP_FIELD = "keysMapBase64"; | ||
private static final String KEY_METADATA_MAP_FIELD = "keysMetadataMap"; | ||
private static final String METADATA_FIELD = "metadata"; | ||
private static final String ENCRYPTION_PARAM_FIELD = "encParamBase64"; | ||
private static final String ALGO_FIELD = "algorithm"; | ||
private static final String COMPRESSION_TYPE_FIELD = "compressionType"; | ||
private static final String UNCPRESSED_MSG_SIZE_FIELD = "uncompressedMessageSize"; | ||
private static final String BATCH_SIZE_FIELD = "batchSize"; | ||
private static final String ENCRYPTION_CTX_FIELD = "encryptionCtx"; | ||
|
||
/** | ||
* Serializes sink-record into json format. It encodes encryption-keys, encryption-param and payload in base64 | ||
* format so, it can be sent in json. | ||
* | ||
* @param inputRecordContext | ||
* @param data | ||
* @return | ||
*/ | ||
public static String serializeRecordToJson(RecordContext inputRecordContext, byte[] data) { | ||
if (inputRecordContext == null) { | ||
return null; | ||
} | ||
JsonObject result = new JsonObject(); | ||
result.addProperty(PAYLOAD_FIELD, getEncoder().encodeToString(data)); | ||
if (inputRecordContext.getProperties() != null) { | ||
JsonObject properties = new JsonObject(); | ||
inputRecordContext.getProperties().entrySet() | ||
.forEach(e -> properties.addProperty(e.getKey(), e.getValue())); | ||
result.add(PROPERTIES_FIELD, properties); | ||
} | ||
if (inputRecordContext.getEncryptionCtx().isPresent()) { | ||
EncryptionContext encryptionCtx = inputRecordContext.getEncryptionCtx().get(); | ||
JsonObject encryptionCtxJson = new JsonObject(); | ||
JsonObject keyBase64Map = new JsonObject(); | ||
JsonObject keyMetadataMap = new JsonObject(); | ||
encryptionCtx.getKeys().entrySet().forEach(entry -> { | ||
keyBase64Map.addProperty(entry.getKey(), getEncoder().encodeToString(entry.getValue().getKeyValue())); | ||
Map<String, String> keyMetadata = entry.getValue().getMetadata(); | ||
if (keyMetadata != null && !keyMetadata.isEmpty()) { | ||
JsonObject metadata = new JsonObject(); | ||
entry.getValue().getMetadata().entrySet() | ||
.forEach(m -> metadata.addProperty(m.getKey(), m.getValue())); | ||
keyMetadataMap.add(entry.getKey(), metadata); | ||
} | ||
}); | ||
encryptionCtxJson.add(KEY_MAP_FIELD, keyBase64Map); | ||
encryptionCtxJson.add(KEY_METADATA_MAP_FIELD, keyMetadataMap); | ||
Map<String, String> metadataMap = encryptionCtx.getMetadata(); | ||
if (metadataMap != null && !metadataMap.isEmpty()) { | ||
JsonObject metadata = new JsonObject(); | ||
encryptionCtx.getMetadata().entrySet().forEach(m -> metadata.addProperty(m.getKey(), m.getValue())); | ||
encryptionCtxJson.add(METADATA_FIELD, metadata); | ||
} | ||
encryptionCtxJson.addProperty(ENCRYPTION_PARAM_FIELD, | ||
getEncoder().encodeToString(encryptionCtx.getParam())); | ||
encryptionCtxJson.addProperty(ALGO_FIELD, encryptionCtx.getAlgorithm()); | ||
if (encryptionCtx.getCompressionType() != null) { | ||
encryptionCtxJson.addProperty(COMPRESSION_TYPE_FIELD, encryptionCtx.getCompressionType().name()); | ||
encryptionCtxJson.addProperty(UNCPRESSED_MSG_SIZE_FIELD, encryptionCtx.getUncompressedMessageSize()); | ||
} | ||
if (encryptionCtx.getBatchSize().isPresent()) { | ||
encryptionCtxJson.addProperty(BATCH_SIZE_FIELD, encryptionCtx.getBatchSize().get()); | ||
} | ||
result.add(ENCRYPTION_CTX_FIELD, encryptionCtxJson); | ||
} | ||
return result.toString(); | ||
} | ||
|
||
} |
162 changes: 162 additions & 0 deletions
162
pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/UtilsTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,162 @@ | ||
/** | ||
* 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.pulsar.io.kinesis; | ||
|
||
import static java.util.Base64.getDecoder; | ||
|
||
import java.util.Map; | ||
import java.util.Optional; | ||
|
||
import org.apache.commons.lang3.StringUtils; | ||
import org.apache.pulsar.common.api.EncryptionContext; | ||
import org.apache.pulsar.common.api.EncryptionContext.EncryptionKey; | ||
import org.apache.pulsar.common.api.proto.PulsarApi.CompressionType; | ||
import org.apache.pulsar.io.core.RecordContext; | ||
import org.testng.Assert; | ||
import org.testng.annotations.Test; | ||
import org.testng.collections.Maps; | ||
|
||
import com.google.gson.Gson; | ||
|
||
import lombok.Getter; | ||
import lombok.Setter; | ||
import lombok.ToString; | ||
|
||
/** | ||
* Unit test of {@link UtilsTest}. | ||
*/ | ||
public class UtilsTest { | ||
|
||
@Test | ||
public void testJsonSerialization() throws Exception { | ||
|
||
final String key1 = "key1"; | ||
final String key2 = "key2"; | ||
final String key1Value = "test1"; | ||
final String key2Value = "test2"; | ||
final String param = "param"; | ||
final String algo = "algo"; | ||
|
||
// prepare encryption-ctx | ||
EncryptionContext ctx = new EncryptionContext(); | ||
ctx.setAlgorithm(algo); | ||
ctx.setBatchSize(Optional.of(10)); | ||
ctx.setCompressionType(CompressionType.LZ4); | ||
ctx.setUncompressedMessageSize(10); | ||
Map<String, EncryptionKey> keys = Maps.newHashMap(); | ||
EncryptionKey encKeyVal = new EncryptionKey(); | ||
encKeyVal.setKeyValue(key1Value.getBytes()); | ||
Map<String, String> metadata1 = Maps.newHashMap(); | ||
metadata1.put("version", "v1"); | ||
metadata1.put("ckms", "cmks-1"); | ||
encKeyVal.setMetadata(metadata1); | ||
EncryptionKey encKeyVal2 = new EncryptionKey(); | ||
encKeyVal2.setKeyValue(key2Value.getBytes()); | ||
Map<String, String> metadata2 = Maps.newHashMap(); | ||
metadata2.put("version", "v2"); | ||
metadata2.put("ckms", "cmks-2"); | ||
encKeyVal2.setMetadata(metadata2); | ||
keys.put(key1, encKeyVal); | ||
keys.put(key2, encKeyVal2); | ||
ctx.setKeys(keys); | ||
ctx.setMetadata(metadata1); | ||
ctx.setParam(param.getBytes()); | ||
|
||
// serialize to json | ||
byte[] data = "payload".getBytes(); | ||
Map<String, String> properties = Maps.newHashMap(); | ||
properties.put("prop1", "value"); | ||
RecordContext recordCtx = new RecordContextImpl(properties, ctx); | ||
String json = Utils.serializeRecordToJson(recordCtx, data); | ||
System.out.println(json); | ||
|
||
// deserialize from json and assert | ||
KinesisMessageResponse kinesisJsonResponse = deSerializeRecordFromJson(json); | ||
Assert.assertEquals(data, getDecoder().decode(kinesisJsonResponse.getPayloadBase64())); | ||
EncryptionCtx encryptionCtxDeser = kinesisJsonResponse.getEncryptionCtx(); | ||
Assert.assertEquals(key1Value.getBytes(), getDecoder().decode(encryptionCtxDeser.getKeysMapBase64().get(key1))); | ||
Assert.assertEquals(key2Value.getBytes(), getDecoder().decode(encryptionCtxDeser.getKeysMapBase64().get(key2))); | ||
Assert.assertEquals(param.getBytes(), getDecoder().decode(encryptionCtxDeser.getEncParamBase64())); | ||
Assert.assertEquals(algo, encryptionCtxDeser.getAlgorithm()); | ||
Assert.assertEquals(metadata1, encryptionCtxDeser.getKeysMetadataMap().get(key1)); | ||
Assert.assertEquals(metadata2, encryptionCtxDeser.getKeysMetadataMap().get(key2)); | ||
Assert.assertEquals(metadata1, encryptionCtxDeser.getMetadata()); | ||
Assert.assertEquals(properties, kinesisJsonResponse.getProperties()); | ||
|
||
} | ||
|
||
class RecordContextImpl implements RecordContext { | ||
Map<String, String> properties; | ||
Optional<EncryptionContext> ectx; | ||
|
||
public RecordContextImpl(Map<String, String> properties, EncryptionContext ectx) { | ||
this.properties = properties; | ||
this.ectx = Optional.of(ectx); | ||
} | ||
|
||
public Map<String, String> getProperties() { | ||
return properties; | ||
} | ||
|
||
public Optional<EncryptionContext> getEncryptionCtx() { | ||
return ectx; | ||
} | ||
} | ||
|
||
public static KinesisMessageResponse deSerializeRecordFromJson(String jsonRecord) { | ||
if (StringUtils.isNotBlank(jsonRecord)) { | ||
return new Gson().fromJson(jsonRecord, KinesisMessageResponse.class); | ||
} | ||
return null; | ||
} | ||
|
||
@ToString | ||
@Setter | ||
@Getter | ||
public static class KinesisMessageResponse { | ||
// Encryption-context if message has been encrypted | ||
private EncryptionCtx encryptionCtx; | ||
// user-properties | ||
private Map<String, String> properties; | ||
// base64 encoded payload | ||
private String payloadBase64; | ||
} | ||
|
||
@ToString | ||
@Setter | ||
@Getter | ||
public static class EncryptionCtx { | ||
// map of encryption-key value. (key-value is base64 encoded) | ||
private Map<String, String> keysMapBase64; | ||
// map of encryption-key metadata | ||
private Map<String, Map<String, String>> keysMetadataMap; | ||
// encryption-ctx metadata | ||
private Map<String, String> metadata; | ||
// encryption param which is base64 encoded | ||
private String encParamBase64; | ||
// encryption algorithm | ||
private String algorithm; | ||
// compression type if message is compressed | ||
private CompressionType compressionType; | ||
private int uncompressedMessageSize; | ||
// number of messages in the batch if msg is batched message | ||
private Integer batchSize; | ||
} | ||
|
||
} |