org.apache.pulsar.client.impl.MessageImpl Java Examples

public static void internalSendAsyncEnter(final Object thiz, final Object arg) {
  if (LocalSpanContext.get(COMPONENT_NAME) != null) {
    LocalSpanContext.get(COMPONENT_NAME).increment();
    return;
  }

  final MessageImpl<?> message = (MessageImpl<?>)arg;
  final Producer<?> producer = (Producer<?>)thiz;

  final Tracer tracer = GlobalTracer.get();
  final Span span = tracer
    .buildSpan("send")
    .withTag(Tags.COMPONENT, COMPONENT_NAME)
    .withTag(Tags.SPAN_KIND, Tags.SPAN_KIND_PRODUCER)
    .withTag(Tags.MESSAGE_BUS_DESTINATION, producer.getTopic())
    .withTag(Tags.PEER_SERVICE, "pulsar")
    .start();

  message.getProperties();

  tracer.inject(span.context(), Builtin.TEXT_MAP, new PropertiesMapInjectAdapter(message.getMessageBuilder()));

  final Scope scope = tracer.activateSpan(span);
  LocalSpanContext.set(COMPONENT_NAME, span, scope);
}
 

@Test
public void peekMessagesTest() throws PulsarAdminException {
    Mockito.when(pulsarAdminService.topics("http://localhost:8080")).thenReturn(topics);
    List<Message<byte[]>> messages = new ArrayList<>();
    messages.add(new MessageImpl<byte[]>("persistent://public/default/test", "1:1", Maps.newTreeMap(), "test".getBytes(), Schema.BYTES));
    Mockito.when(topics.peekMessages("persistent://public/default/test", "sub-1", 1)).thenReturn(messages);

    List<Map<String, Object>> result = topicsService.peekMessages(
            "persistent", "public",
            "default", "test",
            "sub-1", 1,
            "http://localhost:8080");
    Assert.assertEquals(1, result.size());
    result.forEach((message) -> {
        Assert.assertEquals(1L, message.get("ledgerId"));
        Assert.assertEquals(1L, message.get("entryId"));
        Assert.assertEquals(false, message.get("batch"));
        Assert.assertEquals(new String("test".getBytes()), new String((byte[]) message.get("data")));
    });
}
 

/**
 * Delegate work to {@link org.apache.kafka.clients.producer.ProducerInterceptor#onAcknowledgement}
 * @param producer the producer which contains the interceptor.
 * @param message the message that application sends
 * @param msgId the message id that assigned by the broker; null if send failed.
 * @param exception the exception on sending messages, null indicates send has succeed.
 */
@Override
public void onSendAcknowledgement(Producer<byte[]> producer, Message<byte[]> message, MessageId msgId, Throwable exception) {
    try {
        PulsarApi.MessageMetadata.Builder messageMetadataBuilder = ((MessageImpl<byte[]>)message).getMessageBuilder();
        partitionID = getPartitionID(messageMetadataBuilder);
        TopicPartition topicPartition = new TopicPartition(topic, Integer.parseInt(partitionID));
        kafkaProducerInterceptor.onAcknowledgement(new RecordMetadata(topicPartition,
                                                            -1L,
                                                            -1L,
                                                            messageMetadataBuilder.getEventTime(),
                                                            -1L,
                                                            message.getKeyBytes().length,
                                                            message.getValue().length), new Exception(exception));
    } catch (NumberFormatException e) {
        String errorMessage = "Unable to convert partitionID to integer: " + e.getMessage();
        log.error(errorMessage);
        throw new RuntimeException(errorMessage);
    }
}
 

private static PositionImpl findPosition(ReadOnlyCursor readOnlyCursor, long timestamp) throws
        ManagedLedgerException,
        InterruptedException {
    return (PositionImpl) readOnlyCursor.findNewestMatching(SearchAllAvailableEntries, new Predicate<Entry>() {
        @Override
        public boolean apply(Entry entry) {
            MessageImpl msg = null;
            try {
                msg = MessageImpl.deserialize(entry.getDataBuffer());

                return msg.getPublishTime() <= timestamp;
            } catch (Exception e) {
                log.error(e, "Failed To deserialize message when finding position with error: %s", e);
            } finally {
                entry.release();
                if (msg != null) {
                    msg.recycle();
                }
            }
            return false;
        }
    });
}
 

public static ByteBuf messageToByteBuf(Message<byte[]> message) {
    checkArgument(message instanceof MessageImpl);

    MessageImpl<byte[]> msg = (MessageImpl<byte[]>) message;
    MessageMetadata.Builder msgMetadataBuilder = msg.getMessageBuilder();
    ByteBuf payload = msg.getDataBuffer();

    // filled in required fields
    if (!msgMetadataBuilder.hasSequenceId()) {
        msgMetadataBuilder.setSequenceId(-1);
    }
    if (!msgMetadataBuilder.hasPublishTime()) {
        msgMetadataBuilder.setPublishTime(clock.millis());
    }
    if (!msgMetadataBuilder.hasProducerName()) {
        msgMetadataBuilder.setProducerName(FAKE_KOP_PRODUCER_NAME);
    }

    msgMetadataBuilder.setCompression(
        CompressionCodecProvider.convertToWireProtocol(CompressionType.NONE));
    msgMetadataBuilder.setUncompressedSize(payload.readableBytes());
    MessageMetadata msgMetadata = msgMetadataBuilder.build();

    ByteBuf buf = Commands.serializeMetadataAndPayload(ChecksumType.Crc32c, msgMetadata, payload);

    msgMetadataBuilder.recycle();
    msgMetadata.recycle();

    return buf;
}
 

public void findMessages(final long timestamp, AsyncCallbacks.FindEntryCallback callback) {
    this.timestamp = timestamp;
    if (messageFindInProgressUpdater.compareAndSet(this, FALSE, TRUE)) {
        if (log.isDebugEnabled()) {
            log.debug("[{}] Starting message position find at timestamp {}", subName, timestamp);
        }

        cursor.asyncFindNewestMatching(ManagedCursor.FindPositionConstraint.SearchAllAvailableEntries, entry -> {
            MessageImpl msg = null;
            try {
                msg = MessageImpl.deserialize(entry.getDataBuffer());
                return msg.getPublishTime() < timestamp;
            } catch (Exception e) {
                log.error("[{}][{}] Error deserializing message for message position find", topicName, subName, e);
            } finally {
                entry.release();
                if (msg != null) {
                    msg.recycle();
                }
            }
            return false;
        }, this, callback);
    } else {
        if (log.isDebugEnabled()) {
            log.debug("[{}][{}] Ignore message position find scheduled task, last find is still running", topicName,
                    subName);
        }
        callback.findEntryFailed(
                new ManagedLedgerException.ConcurrentFindCursorPositionException("last find is still running"),
                Optional.empty(), null);
    }
}
 

static ProducerSendCallback create(NonPersistentReplicator replicator, Entry entry, MessageImpl msg) {
    ProducerSendCallback sendCallback = RECYCLER.get();
    sendCallback.replicator = replicator;
    sendCallback.entry = entry;
    sendCallback.msg = msg;
    return sendCallback;
}
 

static ProducerSendCallback create(PersistentReplicator replicator, Entry entry, MessageImpl msg) {
    ProducerSendCallback sendCallback = RECYCLER.get();
    sendCallback.replicator = replicator;
    sendCallback.entry = entry;
    sendCallback.msg = msg;
    return sendCallback;
}
 

@Test
public void TestDeleteAction() throws Exception {

    AvroSchema<Foo> schema = AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build());

    Foo deleteObj = new Foo();
    deleteObj.setField3(5);

    byte[] deleteBytes = schema.encode(deleteObj);
    Message<GenericRecord> deleteMessage = mock(MessageImpl.class);
    CompletableFuture<Void> future = new CompletableFuture<>();
    Record<GenericRecord> deleteRecord = PulsarRecord.<GenericRecord>builder()
            .message(deleteMessage)
            .topicName("fake_topic_name")
            .ackFunction(() -> future.complete(null))
            .build();

    GenericSchema<GenericRecord> deleteGenericAvroSchema = new GenericAvroSchema(schema.getSchemaInfo());

    Map<String, String> deleteProperties = Maps.newHashMap();
    deleteProperties.put("ACTION", "DELETE");
    when(deleteMessage.getValue()).thenReturn(deleteGenericAvroSchema.decode(deleteBytes));
    when(deleteMessage.getProperties()).thenReturn(deleteProperties);
    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            deleteObj.toString(),
            deleteMessage.getValue().toString(),
            deleteRecord.getValue().toString());

    jdbcSink.write(deleteRecord);
    future.get(1, TimeUnit.SECONDS);

    // value has been written to db, read it out and verify.
    String deleteQuerySql = "SELECT * FROM " + tableName + " WHERE field3=5";
    Assert.assertEquals(sqliteUtils.select(deleteQuerySql, (resultSet) -> {}), 0);
}
 

private void checkReplicatedSubscriptionMarker(Position position, MessageImpl<?> msg, ByteBuf payload) {
    if (!msg.getMessageBuilder().hasMarkerType()) {
        // No marker is defined
        return;
    }

    int markerType = msg.getMessageBuilder().getMarkerType();

    if (!remoteCluster.equals(msg.getMessageBuilder().getReplicatedFrom())) {
        // Only consider markers that are coming from the same cluster that this
        // replicator instance is assigned to.
        // All the replicators will see all the markers, but we need to only process
        // it once.
        return;
    }

    switch (markerType) {
    case MarkerType.REPLICATED_SUBSCRIPTION_SNAPSHOT_REQUEST_VALUE:
    case MarkerType.REPLICATED_SUBSCRIPTION_SNAPSHOT_RESPONSE_VALUE:
    case MarkerType.REPLICATED_SUBSCRIPTION_UPDATE_VALUE:
        topic.receivedReplicatedSubscriptionMarker(position, markerType, payload);
        break;

    default:
        // Do nothing
    }
}
 

public void expireMessages(int messageTTLInSeconds) {
    if (expirationCheckInProgressUpdater.compareAndSet(this, FALSE, TRUE)) {
        log.info("[{}][{}] Starting message expiry check, ttl= {} seconds", topicName, subName,
                messageTTLInSeconds);

        cursor.asyncFindNewestMatching(ManagedCursor.FindPositionConstraint.SearchActiveEntries, entry -> {
            MessageImpl msg = null;
            try {
                msg = MessageImpl.deserialize(entry.getDataBuffer());
                return msg.isExpired(messageTTLInSeconds);
            } catch (Exception e) {
                log.error("[{}][{}] Error deserializing message for expiry check", topicName, subName, e);
            } finally {
                entry.release();
                if (msg != null) {
                    msg.recycle();
                }
            }
            return false;
        }, this, null);
    } else {
        if (log.isDebugEnabled()) {
            log.debug("[{}][{}] Ignore expire-message scheduled task, last check is still running", topicName,
                    subName);
        }
    }
}
 

private List<Message<byte[]>> getIndividualMsgsFromBatch(String topic, String msgId, byte[] data,
                             Map<String, String> properties, PulsarApi.MessageMetadata.Builder msgMetadataBuilder) {
    List<Message<byte[]>> ret = new ArrayList<>();
    int batchSize = Integer.parseInt(properties.get(BATCH_HEADER));
    ByteBuf buf = Unpooled.wrappedBuffer(data);
    for (int i = 0; i < batchSize; i++) {
        String batchMsgId = msgId + ":" + i;
        PulsarApi.SingleMessageMetadata.Builder singleMessageMetadataBuilder = PulsarApi.SingleMessageMetadata
                .newBuilder();
        try {
            ByteBuf singleMessagePayload =
                    Commands.deSerializeSingleMessageInBatch(buf, singleMessageMetadataBuilder, i, batchSize);
            SingleMessageMetadata singleMessageMetadata = singleMessageMetadataBuilder.build();
            if (singleMessageMetadata.getPropertiesCount() > 0) {
                for (KeyValue entry : singleMessageMetadata.getPropertiesList()) {
                    properties.put(entry.getKey(), entry.getValue());
                }
            }
            ret.add(new MessageImpl<>(topic, batchMsgId, properties, singleMessagePayload,
                    Schema.BYTES, msgMetadataBuilder));
        } catch (Exception ex) {
            log.error("Exception occured while trying to get BatchMsgId: {}", batchMsgId, ex);
        }
        singleMessageMetadataBuilder.recycle();
    }
    buf.release();
    return ret;
}
 

@Override
public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
    MethodInterceptResult result) throws Throwable {
    if (allArguments[0] != null) {
        ProducerEnhanceRequiredInfo requiredInfo = (ProducerEnhanceRequiredInfo) objInst.getSkyWalkingDynamicField();
        ContextCarrier contextCarrier = new ContextCarrier();
        String topicName = requiredInfo.getTopic();
        AbstractSpan activeSpan = ContextManager.createExitSpan(OPERATE_NAME_PREFIX + topicName + PRODUCER_OPERATE_NAME_SUFFIX, contextCarrier, requiredInfo
            .getServiceUrl());
        Tags.MQ_BROKER.set(activeSpan, requiredInfo.getServiceUrl());
        Tags.MQ_TOPIC.set(activeSpan, topicName);
        SpanLayer.asMQ(activeSpan);
        activeSpan.setComponent(ComponentsDefine.PULSAR_PRODUCER);
        CarrierItem next = contextCarrier.items();
        MessageImpl msg = (MessageImpl) allArguments[0];
        while (next.hasNext()) {
            next = next.next();
            msg.getMessageBuilder()
               .addProperties(PulsarApi.KeyValue.newBuilder()
                                                .setKey(next.getHeadKey())
                                                .setValue(next.getHeadValue()));
        }
        if (allArguments.length > 1) {
            EnhancedInstance callbackInstance = (EnhancedInstance) allArguments[1];
            if (callbackInstance != null) {
                ContextSnapshot snapshot = ContextManager.capture();
                if (null != snapshot) {
                    SendCallbackEnhanceRequiredInfo callbackRequiredInfo = new SendCallbackEnhanceRequiredInfo();
                    callbackRequiredInfo.setTopic(topicName);
                    callbackRequiredInfo.setContextSnapshot(snapshot);
                    callbackInstance.setSkyWalkingDynamicField(callbackRequiredInfo);
                }
            }
        }
    }
}
 

@Before
public void setUp() {
    producerInterceptor = new PulsarProducerInterceptor();
    arguments = new Object[] {
        msg,
        null
    };
    argumentType = new Class[] {MessageImpl.class};
}
 

public static MessageImpl<byte[]> recordToEntry(Record record) {
    @SuppressWarnings("unchecked")
    TypedMessageBuilderImpl<byte[]> builder = new TypedMessageBuilderImpl(null, Schema.BYTES);

    // key
    if (record.hasKey()) {
        byte[] key = new byte[record.keySize()];
        record.key().get(key);
        builder.keyBytes(key);
        // reuse ordering key to avoid converting string < > bytes
        builder.orderingKey(key);
    }

    // value
    if (record.hasValue()) {
        byte[] value = new byte[record.valueSize()];
        record.value().get(value);
        builder.value(value);
    } else {
        builder.value(new byte[0]);
    }

    // sequence
    if (record.sequence() >= 0) {
        builder.sequenceId(record.sequence());
    }

    // timestamp
    if (record.timestamp() >= 0) {
        builder.eventTime(record.timestamp());
    }

    // header
    for (Header h : record.headers()) {
        builder.property(h.key(),
            new String(h.value(), UTF_8));
    }

    return (MessageImpl<byte[]>) builder.getMessage();
}
 

public void findMessages(final long timestamp, AsyncCallbacks.FindEntryCallback callback) {
    this.timestamp = timestamp;
    if (messageFindInProgressUpdater.compareAndSet(this, FALSE, TRUE)) {
        if (log.isDebugEnabled()) {
            log.debug("[{}] Starting message position find at timestamp {}",  timestamp);
        }

        asyncFindNewestMatching(ManagedCursor.FindPositionConstraint.SearchAllAvailableEntries, entry -> {
            MessageImpl msg = null;
            try {
                msg = MessageImpl.deserialize(entry.getDataBuffer());
                return msg.getPublishTime() <= timestamp;
            } catch (Exception e) {
                log.error("[{}][{}] Error deserializing message for message position find",  e);
            } finally {
                entry.release();
                if (msg != null) {
                    msg.recycle();
                }
            }
            return false;
        }, this, callback);
    } else {
        if (log.isDebugEnabled()) {
            log.debug("[{}][{}] Ignore message position find scheduled task, last find is still running");
        }
        callback.findEntryFailed(
            new ManagedLedgerException.ConcurrentFindCursorPositionException("last find is still running"),
            Optional.empty(),
            null);
    }
}
 

private List<Message<byte[]>> getMessagesFromHttpResponse(String topic, Response response) throws Exception {

        if (response.getStatus() != Status.OK.getStatusCode()) {
            throw getApiException(response);
        }

        String msgId = response.getHeaderString(MESSAGE_ID);
        PulsarApi.MessageMetadata.Builder messageMetadata = PulsarApi.MessageMetadata.newBuilder();
        try (InputStream stream = (InputStream) response.getEntity()) {
            byte[] data = new byte[stream.available()];
            stream.read(data);

            Map<String, String> properties = Maps.newTreeMap();
            MultivaluedMap<String, Object> headers = response.getHeaders();
            Object tmp = headers.getFirst(PUBLISH_TIME);
            if (tmp != null) {
                properties.put("publish-time", (String) tmp);
            }

            tmp = headers.getFirst("X-Pulsar-null-value");
            if (tmp != null) {
                messageMetadata.setNullValue(Boolean.parseBoolean(tmp.toString()));
            }

            tmp = headers.getFirst(BATCH_HEADER);
            if (response.getHeaderString(BATCH_HEADER) != null) {
                properties.put(BATCH_HEADER, (String) tmp);
                return getIndividualMsgsFromBatch(topic, msgId, data, properties, messageMetadata);
            }
            for (Entry<String, List<Object>> entry : headers.entrySet()) {
                String header = entry.getKey();
                if (header.contains("X-Pulsar-PROPERTY-")) {
                    String keyName = header.substring("X-Pulsar-PROPERTY-".length());
                    properties.put(keyName, (String) entry.getValue().get(0));
                }
            }

            return Collections.singletonList(new MessageImpl<byte[]>(topic, msgId, properties,
                    Unpooled.wrappedBuffer(data), Schema.BYTES, messageMetadata));
        }
    }
 

@Override
public MessageImpl<?> getNextMessage() {
    return null;
}
 

@Override
public void addCallback(MessageImpl<?> msg, SendCallback scb) {
    // noop
}
 

public static Message getPulsarMessage(String messageId, byte[] payload) {
  return new MessageImpl<byte[]>(messageId, messageId, new HashMap<>(), Unpooled.wrappedBuffer(payload), null);
}
 

@Override
public MessageImpl<?> getNextMessage() {
    return null;
}
 

@Override
public void addCallback(MessageImpl<?> msg, SendCallback scb) {
    // noop
}
 

public void sendMessage(Entry entry) {
    if ((STATE_UPDATER.get(this) == State.Started) && isWritable()) {

        int length = entry.getLength();
        ByteBuf headersAndPayload = entry.getDataBuffer();
        MessageImpl msg;
        try {
            msg = MessageImpl.deserialize(headersAndPayload);
        } catch (Throwable t) {
            log.error("[{}][{} -> {}] Failed to deserialize message at {} (buffer size: {}): {}", topicName,
                    localCluster, remoteCluster, entry.getPosition(), length, t.getMessage(), t);
            entry.release();
            return;
        }

        if (msg.isReplicated()) {
            // Discard messages that were already replicated into this region
            entry.release();
            msg.recycle();
            return;
        }

        if (msg.hasReplicateTo() && !msg.getReplicateTo().contains(remoteCluster)) {
            if (log.isDebugEnabled()) {
                log.debug("[{}][{} -> {}] Skipping message at {} / msg-id: {}: replicateTo {}", topicName,
                        localCluster, remoteCluster, entry.getPosition(), msg.getMessageId(), msg.getReplicateTo());
            }
            entry.release();
            msg.recycle();
            return;
        }

        msgOut.recordEvent(headersAndPayload.readableBytes());

        msg.setReplicatedFrom(localCluster);

        headersAndPayload.retain();

        producer.sendAsync(msg, ProducerSendCallback.create(this, entry, msg));

    } else {
        if (log.isDebugEnabled()) {
            log.debug("[{}][{} -> {}] dropping message because replicator producer is not started/writable",
                    topicName, localCluster, remoteCluster);
        }
        msgDrop.recordEvent();
        entry.release();
    }
}
 

@Test
public void TestOpenAndWriteSink() throws Exception {
    message = mock(MessageImpl.class);
    Map<String, Object> configs = new HashMap<>();
    configs.put("solrUrl", "http://localhost:8983/solr");
    configs.put("solrMode", "Standalone");
    configs.put("solrCollection", "techproducts");
    configs.put("solrCommitWithinMs", "100");
    configs.put("username", "");
    configs.put("password", "");
    GenericSchema<GenericRecord> genericAvroSchema;

    SolrGenericRecordSink sink = new SolrGenericRecordSink();

    // prepare a foo Record
    Foo obj = new Foo();
    obj.setField1("FakeFiled1");
    obj.setField2("FakeFiled1");
    AvroSchema<Foo> schema = AvroSchema.of(Foo.class);

    byte[] bytes = schema.encode(obj);
    AutoConsumeSchema autoConsumeSchema = new AutoConsumeSchema();
    autoConsumeSchema.setSchema(GenericSchemaImpl.of(schema.getSchemaInfo()));

    Record<GenericRecord> record = PulsarRecord.<GenericRecord>builder()
        .message(message)
        .topicName("fake_topic_name")
        .build();

    genericAvroSchema = new GenericAvroSchema(schema.getSchemaInfo());

    when(message.getValue())
            .thenReturn(genericAvroSchema.decode(bytes));

    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
        obj.toString(),
        message.getValue().toString(),
        record.getValue().toString());

    // open should success
    sink.open(configs, null);
}
 

@Test(enabled = false)
public void TestOpenAndWriteSink() throws Exception {
    message = mock(MessageImpl.class);
    GenericSchema<GenericRecord> genericAvroSchema;

    Map<String, Object> map = new HashMap<>();
    map.put("zookeeperQuorum", "localhost");
    map.put("zookeeperClientPort", "2181");
    map.put("zookeeperZnodeParent", "/hbase");
    map.put("tableName", "default:pulsar_hbase");
    map.put("rowKeyName", rowKeyName);
    map.put("familyName", familyName);

    List<String> qualifierNames = new ArrayList<>();
    qualifierNames.add(name);
    qualifierNames.add(address);
    qualifierNames.add(age);
    qualifierNames.add(flag);
    map.put("qualifierNames",qualifierNames);

    mockSinkContext = mock(SinkContext.class);
    HbaseGenericRecordSink sink = new HbaseGenericRecordSink();

    // prepare a foo Record
    Foo obj = new Foo();
    obj.setRowKey("rowKey_value");
    obj.setName("name_value");
    obj.setAddress("address_value");
    obj.setAge(30);
    obj.setFlag(true);
    AvroSchema<Foo> schema = AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build());

    byte[] bytes = schema.encode(obj);
    AutoConsumeSchema autoConsumeSchema = new AutoConsumeSchema();
    autoConsumeSchema.setSchema(GenericSchemaImpl.of(schema.getSchemaInfo()));

    PulsarSourceConfig pulsarSourceConfig = new PulsarSourceConfig();
    Consumer consumer = mock(Consumer.class);

    Record<GenericRecord> record = PulsarRecord.<GenericRecord>builder()
        .message(message)
        .topicName("fake_topic_name")
        .ackFunction(() -> {
            if (pulsarSourceConfig
                    .getProcessingGuarantees() == FunctionConfig.ProcessingGuarantees.EFFECTIVELY_ONCE) {
                consumer.acknowledgeCumulativeAsync(message);
            } else {
                consumer.acknowledgeAsync(message);
            }
        }).failFunction(() -> {
            if (pulsarSourceConfig.getProcessingGuarantees() == FunctionConfig.ProcessingGuarantees.EFFECTIVELY_ONCE) {
                throw new RuntimeException("Failed to process message: " + message.getMessageId());
            }
        })
        .build();

    genericAvroSchema = new GenericAvroSchema(schema.getSchemaInfo());

    when(message.getValue())
            .thenReturn(genericAvroSchema.decode(bytes));

    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            obj.toString(),
            message.getValue().toString(),
            record.getValue().toString());

    // change batchSize to 1, to flush on each write.
    map.put("batchTimeMs", 1);
    map.put("batchSize", 1);
    // open should success
    sink.open(map,mockSinkContext);

    // write should success.
    sink.write(record);
    log.info("executed write");
    // sleep to wait backend flush complete
    Thread.sleep(500);

    // value has been written to hbase table, read it out and verify.
    Table table = TableUtils.getTable(map);
    Get scan = new Get(Bytes.toBytes(obj.getRowKey()));
    Result result = table.get(scan);
    byte[] byteName = result.getValue(Bytes.toBytes(familyName), Bytes.toBytes(name));
    byte[] byteAddress = result.getValue(Bytes.toBytes(familyName), Bytes.toBytes(address));
    byte[] byteAge = result.getValue(Bytes.toBytes(familyName), Bytes.toBytes(age));
    byte[] byteFlag = result.getValue(Bytes.toBytes(familyName), Bytes.toBytes(flag));
    Assert.assertEquals(obj.getName(), Bytes.toString(byteName));
    Assert.assertEquals(obj.getAddress(), Bytes.toString(byteAddress));
    Assert.assertEquals(obj.getAge(), Bytes.toInt(byteAge));
    Assert.assertEquals(obj.isFlag(), Bytes.toBoolean(byteFlag));

    table.close();
    sink.close();
}
 

@Test
public void TestUpdateAction() throws Exception {

    AvroSchema<Foo> schema = AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build());

    Foo updateObj = new Foo();
    updateObj.setField1("ValueOfField3");
    updateObj.setField2("ValueOfField3");
    updateObj.setField3(4);

    byte[] updateBytes = schema.encode(updateObj);
    Message<GenericRecord> updateMessage = mock(MessageImpl.class);
    CompletableFuture<Void> future = new CompletableFuture<>();
    Record<GenericRecord> updateRecord = PulsarRecord.<GenericRecord>builder()
            .message(updateMessage)
            .topicName("fake_topic_name")
            .ackFunction(() -> future.complete(null))
            .build();

    GenericSchema<GenericRecord> updateGenericAvroSchema;
    updateGenericAvroSchema = new GenericAvroSchema(schema.getSchemaInfo());

    Map<String, String> updateProperties = Maps.newHashMap();
    updateProperties.put("ACTION", "UPDATE");
    when(updateMessage.getValue()).thenReturn(updateGenericAvroSchema.decode(updateBytes));
    when(updateMessage.getProperties()).thenReturn(updateProperties);
    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            updateObj.toString(),
            updateMessage.getValue().toString(),
            updateRecord.getValue().toString());

    jdbcSink.write(updateRecord);
    future.get(1, TimeUnit.SECONDS);

    // value has been written to db, read it out and verify.
    String updateQuerySql = "SELECT * FROM " + tableName + " WHERE field3=4";
    sqliteUtils.select(updateQuerySql, (resultSet) -> {
        Assert.assertEquals(updateObj.getField1(), resultSet.getString(1));
        Assert.assertEquals(updateObj.getField2(), resultSet.getString(2));
        Assert.assertEquals(updateObj.getField3(), resultSet.getInt(3));
    });
}
 

private void testOpenAndWriteSink(Map<String, String> actionProperties) throws Exception {
    Message<GenericRecord> insertMessage = mock(MessageImpl.class);
    GenericSchema<GenericRecord> genericAvroSchema;
    // prepare a foo Record
    Foo insertObj = new Foo();
    insertObj.setField1("ValueOfField1");
    insertObj.setField2("ValueOfField2");
    insertObj.setField3(3);
    AvroSchema<Foo> schema = AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build());

    byte[] insertBytes = schema.encode(insertObj);
    CompletableFuture<Void> future = new CompletableFuture<>();
    Record<GenericRecord> insertRecord = PulsarRecord.<GenericRecord>builder()
        .message(insertMessage)
        .topicName("fake_topic_name")
        .ackFunction(() -> future.complete(null))
        .build();

    genericAvroSchema = new GenericAvroSchema(schema.getSchemaInfo());
    when(insertMessage.getValue()).thenReturn(genericAvroSchema.decode(insertBytes));
    when(insertMessage.getProperties()).thenReturn(actionProperties);
    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            insertObj.toString(),
            insertMessage.getValue().toString(),
            insertRecord.getValue().toString());

    // write should success.
    jdbcSink.write(insertRecord);
    log.info("executed write");
    // sleep to wait backend flush complete
    future.get(1, TimeUnit.SECONDS);

    // value has been written to db, read it out and verify.
    String querySql = "SELECT * FROM " + tableName + " WHERE field3=3";
    int count = sqliteUtils.select(querySql, (resultSet) -> {
        Assert.assertEquals(insertObj.getField1(), resultSet.getString(1));
        Assert.assertEquals(insertObj.getField2(), resultSet.getString(2));
        Assert.assertEquals(insertObj.getField3(), resultSet.getInt(3));
    });
    Assert.assertEquals(count, 1);

}
 

private void testOpenAndWriteSinkNullValueJson(Map<String, String> actionProperties) throws Exception {
    Message<GenericRecord> insertMessage = mock(MessageImpl.class);
    GenericSchema<GenericRecord> genericAvroSchema;
    // prepare a foo Record
    Foo insertObj = new Foo();
    insertObj.setField1("ValueOfField1");
    // Not setting field2
    // Field1 is the key and field3 is used for selecting records 
    insertObj.setField3(3);
    JSONSchema<Foo> schema = JSONSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).withAlwaysAllowNull(true).build());

    byte[] insertBytes = schema.encode(insertObj);
    CompletableFuture<Void> future = new CompletableFuture<>();
    Record<GenericRecord> insertRecord = PulsarRecord.<GenericRecord>builder()
        .message(insertMessage)
        .topicName("fake_topic_name")
        .ackFunction(() -> future.complete(null))
        .build();

    GenericSchema<GenericRecord> decodeSchema = GenericSchemaImpl.of(schema.getSchemaInfo());
    when(insertMessage.getValue()).thenReturn(decodeSchema.decode(insertBytes));
    when(insertMessage.getProperties()).thenReturn(actionProperties);
    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            insertObj.toString(),
            insertMessage.getValue().toString(),
            insertRecord.getValue().toString());

    // write should success.
    jdbcSink.write(insertRecord);
    log.info("executed write");
    // sleep to wait backend flush complete
    future.get(1, TimeUnit.SECONDS);

    // value has been written to db, read it out and verify.
    String querySql = "SELECT * FROM " + tableName + " WHERE field3=3";
    int count = sqliteUtils.select(querySql, (resultSet) -> {
        Assert.assertEquals(insertObj.getField1(), resultSet.getString(1));
        Assert.assertNull(insertObj.getField2());
        Assert.assertEquals(insertObj.getField3(), resultSet.getInt(3));
    });
    Assert.assertEquals(count, 1);

}
 

private void testOpenAndWriteSinkJson(Map<String, String> actionProperties) throws Exception {
    Message<GenericRecord> insertMessage = mock(MessageImpl.class);
    GenericSchema<GenericRecord> genericAvroSchema;
    // prepare a foo Record
    Foo insertObj = new Foo();
    insertObj.setField1("ValueOfField1");
    insertObj.setField2("ValueOfField2");
    insertObj.setField3(3);
    JSONSchema<Foo> schema = JSONSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).withAlwaysAllowNull(true).build());

    byte[] insertBytes = schema.encode(insertObj);
    CompletableFuture<Void> future = new CompletableFuture<>();
    Record<GenericRecord> insertRecord = PulsarRecord.<GenericRecord>builder()
        .message(insertMessage)
        .topicName("fake_topic_name")
        .ackFunction(() -> future.complete(null))
        .build();

    GenericSchema<GenericRecord> decodeSchema = GenericSchemaImpl.of(schema.getSchemaInfo());
    when(insertMessage.getValue()).thenReturn(decodeSchema.decode(insertBytes));
    when(insertMessage.getProperties()).thenReturn(actionProperties);
    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            insertObj.toString(),
            insertMessage.getValue().toString(),
            insertRecord.getValue().toString());

    // write should success.
    jdbcSink.write(insertRecord);
    log.info("executed write");
    // sleep to wait backend flush complete
    future.get(1, TimeUnit.SECONDS);

    // value has been written to db, read it out and verify.
    String querySql = "SELECT * FROM " + tableName + " WHERE field3=3";
    int count = sqliteUtils.select(querySql, (resultSet) -> {
        Assert.assertEquals(insertObj.getField1(), resultSet.getString(1));
        Assert.assertEquals(insertObj.getField2(), resultSet.getString(2));
        Assert.assertEquals(insertObj.getField3(), resultSet.getInt(3));
    });
    Assert.assertEquals(count, 1);

}
 

private void testOpenAndWriteSinkNullValue(Map<String, String> actionProperties) throws Exception {
    Message<GenericRecord> insertMessage = mock(MessageImpl.class);
    GenericSchema<GenericRecord> genericAvroSchema;
    // prepare a foo Record
    Foo insertObj = new Foo();
    insertObj.setField1("ValueOfField1");
    // Not setting field2
    // Field1 is the key and field3 is used for selecting records 
    insertObj.setField3(3);
    AvroSchema<Foo> schema = AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).withAlwaysAllowNull(true).build());

    byte[] insertBytes = schema.encode(insertObj);
    CompletableFuture<Void> future = new CompletableFuture<>();
    Record<GenericRecord> insertRecord = PulsarRecord.<GenericRecord>builder()
        .message(insertMessage)
        .topicName("fake_topic_name")
        .ackFunction(() -> future.complete(null))
        .build();

    genericAvroSchema = new GenericAvroSchema(schema.getSchemaInfo());
    when(insertMessage.getValue()).thenReturn(genericAvroSchema.decode(insertBytes));
    when(insertMessage.getProperties()).thenReturn(actionProperties);
    log.info("foo:{}, Message.getValue: {}, record.getValue: {}",
            insertObj.toString(),
            insertMessage.getValue().toString(),
            insertRecord.getValue().toString());

    // write should success.
    jdbcSink.write(insertRecord);
    log.info("executed write");
    // sleep to wait backend flush complete
    future.get(1, TimeUnit.SECONDS);

    // value has been written to db, read it out and verify.
    String querySql = "SELECT * FROM " + tableName + " WHERE field3=3";
    int count = sqliteUtils.select(querySql, (resultSet) -> {
        Assert.assertEquals(insertObj.getField1(), resultSet.getString(1));
        Assert.assertNull(insertObj.getField2());
        Assert.assertEquals(insertObj.getField3(), resultSet.getInt(3));
    });
    Assert.assertEquals(count, 1);

}