Add PartialSerialisedMessage test. #27452
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -109,10 +109,9 @@ public void accept(final String messageString, final Integer sizeInBytes) throws | |
| */ | ||
| deserializeAirbyteMessage(messageString) | ||
| .ifPresent(message -> { | ||
| if (message.getType().equals(Type.RECORD)) { | ||
| validateRecord(message); | ||
| } | ||
| bufferEnqueue.addRecord(message, sizeInBytes + PARTIAL_DESERIALIZE_REF_BYTES); | ||
| }); | ||
| } | ||
|
|
@@ -121,18 +120,28 @@ public void accept(final String messageString, final Integer sizeInBytes) throws | |
| * Deserializes to a {@link PartialAirbyteMessage} which can represent both a Record or a State | ||
| * Message | ||
| * | ||
| * PartialAirbyteMessage holds either: | ||
| * <li>entire serialized message string when message is a valid State Message | ||
| * <li>serialized AirbyteRecordMessage when message is a valid Record Message</li> | ||
| * | ||
| * @param messageString the string to deserialize | ||
| * @return PartialAirbyteMessage if the message is valid, empty otherwise | ||
| */ | ||
| @VisibleForTesting | ||
| public static Optional<PartialAirbyteMessage> deserializeAirbyteMessage(final String messageString) { | ||
| // TODO: (ryankfu) plumb in the serialized AirbyteStateMessage to match AirbyteRecordMessage and | ||
| // code parity | ||
| final Optional<PartialAirbyteMessage> messageOptional = Jsons.tryDeserialize(messageString, PartialAirbyteMessage.class) | ||
| .map(partial -> { | ||
| if (partial.getType().equals(Type.RECORD) && partial.getRecord().getData() != null) { | ||
| return partial.withSerialized(partial.getRecord().getData().toString()); | ||
| } else if (partial.getType().equals(Type.STATE)) { | ||
| return partial.withSerialized(messageString); | ||
| } else { | ||
|
There was a problem hiding this comment. @davinchia added another case where I believe we can have bad data which is when messageString is neither a STATE or RECORD message. In this case, this should not be passed to the consumer. I haven't personally seen this happen but this now more closely matches the description that PartialAirbyteMessage should only contain either STATE or RECORD. Lmk if you think this is unnecessary. Not sure if the platform filters out messages There was a problem hiding this comment. We cannot guarantee messages arrive at the Destination ungarbled, so we definitely need to consider this case. Less decided whether we should fail fast or log an error.. I think we should do the simplest for now and throw a RTE to make sure we don't drop data. I'm wondering if this happens in practice somehow. We should add logging + test this out before committing this code. There was a problem hiding this comment. Sounds good, I'll add an RTE here, add some logging, publish a new pre-release version, and run it overnight with some resets to see if this ever gets triggered but agree we should fail fast here |
||
| return null; | ||
| } | ||
| }); | ||
|
|
||
| if (messageOptional.isPresent()) { | ||
| return messageOptional; | ||
| } else { | ||
|
|
@@ -143,6 +152,7 @@ private Optional<PartialAirbyteMessage> deserializeAirbyteMessage(final String m | |
| return Optional.empty(); | ||
| } | ||
| } | ||
|
|
||
| } | ||
|
|
||
| /** | ||
|
|
@@ -151,7 +161,8 @@ private Optional<PartialAirbyteMessage> deserializeAirbyteMessage(final String m | |
| * @param input a JSON string that represents an {@link AirbyteMessage}. | ||
| * @return {@code true} if the message is a state message, {@code false} otherwise. | ||
| */ | ||
| @VisibleForTesting | ||
| public static boolean isStateMessage(final String input) { | ||
| final Optional<AirbyteTypeMessage> deserialized = Jsons.tryDeserialize(input, AirbyteTypeMessage.class); | ||
| return deserialized.filter(airbyteTypeMessage -> airbyteTypeMessage.getType() == Type.STATE).isPresent(); | ||
| } | ||
|
|
||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,66 @@ | ||
| /* | ||
| * Copyright (c) 2023 Airbyte, Inc., all rights reserved. | ||
| */ | ||
|
|
||
| package io.airbyte.integrations.destination_async; | ||
|
|
||
| import io.airbyte.commons.json.Jsons; | ||
| import io.airbyte.integrations.destination_async.partial_messages.PartialAirbyteMessage; | ||
| import io.airbyte.protocol.models.AirbyteStateMessage; | ||
| import io.airbyte.protocol.models.AirbyteStreamState; | ||
| import io.airbyte.protocol.models.StreamDescriptor; | ||
| import io.airbyte.protocol.models.v0.AirbyteMessage; | ||
| import io.airbyte.protocol.models.v0.AirbyteRecordMessage; | ||
| import java.time.Instant; | ||
| import org.junit.jupiter.api.Assertions; | ||
| import org.junit.jupiter.api.Test; | ||
|
|
||
| public class PartialAirbyteMessageTest { | ||
|
|
||
| @Test | ||
| void testDeserializeRecord() { | ||
| final long emittedAt = Instant.now().toEpochMilli(); | ||
| final var serializedRec = Jsons.serialize(new AirbyteMessage() | ||
| .withType(AirbyteMessage.Type.RECORD) | ||
| .withRecord(new AirbyteRecordMessage() | ||
| .withStream("users") | ||
| .withNamespace("public") | ||
| .withEmittedAt(emittedAt) | ||
| .withData(Jsons.jsonNode("data")))); | ||
|
|
||
| final var rec = Jsons.tryDeserialize(serializedRec, PartialAirbyteMessage.class).get(); | ||
| Assertions.assertEquals(AirbyteMessage.Type.RECORD, rec.getType()); | ||
| Assertions.assertEquals("users", rec.getRecord().getStream()); | ||
| Assertions.assertEquals("public", rec.getRecord().getNamespace()); | ||
| Assertions.assertEquals("\"data\"", rec.getRecord().getData().toString()); | ||
| Assertions.assertEquals(emittedAt, rec.getRecord().getEmittedAt()); | ||
| } | ||
|
|
||
| @Test | ||
| void testDeserializeState() { | ||
| final var serializedState = Jsons.serialize(new io.airbyte.protocol.models.AirbyteMessage() | ||
| .withType(io.airbyte.protocol.models.AirbyteMessage.Type.STATE) | ||
| .withState(new AirbyteStateMessage().withStream( | ||
| new AirbyteStreamState().withStreamDescriptor( | ||
| new StreamDescriptor().withName("user").withNamespace("public")) | ||
| .withStreamState(Jsons.jsonNode("data"))) | ||
| .withType(AirbyteStateMessage.AirbyteStateType.STREAM))); | ||
|
|
||
| final var rec = Jsons.tryDeserialize(serializedState, PartialAirbyteMessage.class).get(); | ||
| Assertions.assertEquals(AirbyteMessage.Type.STATE, rec.getType()); | ||
|
|
||
| final var streamDesc = rec.getState().getStream().getStreamDescriptor(); | ||
| Assertions.assertEquals("user", streamDesc.getName()); | ||
| Assertions.assertEquals("public", streamDesc.getNamespace()); | ||
| Assertions.assertEquals(io.airbyte.protocol.models.v0.AirbyteStateMessage.AirbyteStateType.STREAM, rec.getState().getType()); | ||
| } | ||
|
|
||
| @Test | ||
| void testGarbage() { | ||
| final var badSerialization = "messed up data"; | ||
|
|
||
| final var rec = Jsons.tryDeserialize(badSerialization, PartialAirbyteMessage.class); | ||
| Assertions.assertTrue(rec.isEmpty()); | ||
| } | ||
|
|
||
| } |
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
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we should add tests for this since this is now a critical piece