-
Notifications
You must be signed in to change notification settings - Fork 4.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Update state.state type * Add state aggregator * Test and format * PR comments * Move to its own package * Update airbyte-workers/src/test/java/io/airbyte/workers/internal/state_aggregator/StateAggregatorTest.java Co-authored-by: Lake Mossman <lake@airbyte.io> * format * Update airbyte-workers/src/main/java/io/airbyte/workers/internal/state_aggregator/DefaultStateAggregator.java Co-authored-by: Lake Mossman <lake@airbyte.io> * format Co-authored-by: Lake Mossman <lake@airbyte.io>
- Loading branch information
1 parent
ca272c3
commit 59e20f2
Showing
7 changed files
with
313 additions
and
5 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
55 changes: 55 additions & 0 deletions
55
...rs/src/main/java/io/airbyte/workers/internal/state_aggregator/DefaultStateAggregator.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,55 @@ | ||
/* | ||
* Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.workers.internal.state_aggregator; | ||
|
||
import com.google.common.base.Preconditions; | ||
import io.airbyte.config.State; | ||
import io.airbyte.protocol.models.AirbyteStateMessage; | ||
import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType; | ||
|
||
public class DefaultStateAggregator implements StateAggregator { | ||
|
||
private AirbyteStateType stateType = null; | ||
private final StateAggregator streamStateAggregator = new StreamStateAggregator(); | ||
private final StateAggregator singleStateAggregator = new SingleStateAggregator(); | ||
|
||
@Override | ||
public void ingest(final AirbyteStateMessage stateMessage) { | ||
checkTypeOrSetType(stateMessage.getType()); | ||
|
||
getStateAggregator().ingest(stateMessage); | ||
} | ||
|
||
@Override | ||
public State getAggregated() { | ||
return getStateAggregator().getAggregated(); | ||
} | ||
|
||
/** | ||
* Return the state aggregator that match the state type. | ||
*/ | ||
private StateAggregator getStateAggregator() { | ||
return switch (stateType) { | ||
case STREAM -> streamStateAggregator; | ||
case GLOBAL, LEGACY -> singleStateAggregator; | ||
}; | ||
} | ||
|
||
/** | ||
* We can not have 2 different state types given to the same instance of this class. This method set | ||
* the type if it is not. If the state type doesn't exist in the message, it is set to LEGACY | ||
*/ | ||
private void checkTypeOrSetType(AirbyteStateType inputStateType) { | ||
if (inputStateType == null) { | ||
inputStateType = AirbyteStateType.LEGACY; | ||
} | ||
if (this.stateType == null) { | ||
this.stateType = inputStateType; | ||
} | ||
Preconditions.checkArgument(this.stateType == inputStateType, | ||
"Input state type " + inputStateType + " does not match the aggregator's current state type " + this.stateType); | ||
} | ||
|
||
} |
32 changes: 32 additions & 0 deletions
32
...ers/src/main/java/io/airbyte/workers/internal/state_aggregator/SingleStateAggregator.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,32 @@ | ||
/* | ||
* Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.workers.internal.state_aggregator; | ||
|
||
import io.airbyte.commons.json.Jsons; | ||
import io.airbyte.config.State; | ||
import io.airbyte.protocol.models.AirbyteStateMessage; | ||
import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType; | ||
import java.util.List; | ||
|
||
class SingleStateAggregator implements StateAggregator { | ||
|
||
AirbyteStateMessage state; | ||
|
||
@Override | ||
public void ingest(final AirbyteStateMessage stateMessage) { | ||
state = stateMessage; | ||
} | ||
|
||
@Override | ||
public State getAggregated() { | ||
if (state.getType() == null || state.getType() == AirbyteStateType.LEGACY) { | ||
return new State().withState(state.getData()); | ||
} else { | ||
return new State() | ||
.withState(Jsons.jsonNode(List.of(state))); | ||
} | ||
} | ||
|
||
} |
16 changes: 16 additions & 0 deletions
16
...e-workers/src/main/java/io/airbyte/workers/internal/state_aggregator/StateAggregator.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,16 @@ | ||
/* | ||
* Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.workers.internal.state_aggregator; | ||
|
||
import io.airbyte.config.State; | ||
import io.airbyte.protocol.models.AirbyteStateMessage; | ||
|
||
public interface StateAggregator { | ||
|
||
void ingest(AirbyteStateMessage stateMessage); | ||
|
||
State getAggregated(); | ||
|
||
} |
31 changes: 31 additions & 0 deletions
31
...ers/src/main/java/io/airbyte/workers/internal/state_aggregator/StreamStateAggregator.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,31 @@ | ||
/* | ||
* Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.workers.internal.state_aggregator; | ||
|
||
import io.airbyte.commons.json.Jsons; | ||
import io.airbyte.config.State; | ||
import io.airbyte.protocol.models.AirbyteStateMessage; | ||
import io.airbyte.protocol.models.StreamDescriptor; | ||
import java.util.HashMap; | ||
import java.util.Map; | ||
|
||
class StreamStateAggregator implements StateAggregator { | ||
|
||
Map<StreamDescriptor, AirbyteStateMessage> aggregatedState = new HashMap<>(); | ||
|
||
@Override | ||
public void ingest(final AirbyteStateMessage stateMessage) { | ||
aggregatedState.put(stateMessage.getStream().getStreamDescriptor(), stateMessage); | ||
} | ||
|
||
@Override | ||
public State getAggregated() { | ||
|
||
return new State() | ||
.withState( | ||
Jsons.jsonNode(aggregatedState.values())); | ||
} | ||
|
||
} |
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
162 changes: 162 additions & 0 deletions
162
...rkers/src/test/java/io/airbyte/workers/internal/state_aggregator/StateAggregatorTest.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 @@ | ||
/* | ||
* Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.workers.internal.state_aggregator; | ||
|
||
import static io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType.GLOBAL; | ||
import static io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType.LEGACY; | ||
import static io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType.STREAM; | ||
|
||
import com.google.common.collect.Lists; | ||
import io.airbyte.commons.json.Jsons; | ||
import io.airbyte.config.State; | ||
import io.airbyte.protocol.models.AirbyteGlobalState; | ||
import io.airbyte.protocol.models.AirbyteStateMessage; | ||
import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType; | ||
import io.airbyte.protocol.models.AirbyteStreamState; | ||
import io.airbyte.protocol.models.StreamDescriptor; | ||
import java.util.Arrays; | ||
import java.util.List; | ||
import java.util.stream.Stream; | ||
import org.assertj.core.api.Assertions; | ||
import org.junit.jupiter.api.BeforeEach; | ||
import org.junit.jupiter.api.Test; | ||
import org.junit.jupiter.params.ParameterizedTest; | ||
import org.junit.jupiter.params.provider.EnumSource; | ||
|
||
public class StateAggregatorTest { | ||
|
||
StateAggregator stateAggregator; | ||
|
||
@BeforeEach | ||
public void init() { | ||
stateAggregator = new DefaultStateAggregator(); | ||
} | ||
|
||
@ParameterizedTest | ||
@EnumSource(AirbyteStateType.class) | ||
public void testCantMixType(final AirbyteStateType stateType) { | ||
final Stream<AirbyteStateType> allTypes = Arrays.stream(AirbyteStateType.values()); | ||
|
||
stateAggregator.ingest(getEmptyMessage(stateType)); | ||
|
||
final List<AirbyteStateType> differentTypes = allTypes.filter(type -> type != stateType).toList(); | ||
differentTypes.forEach(differentType -> Assertions.assertThatThrownBy(() -> stateAggregator.ingest(getEmptyMessage(differentType)))); | ||
} | ||
|
||
@Test | ||
public void testCantMixNullType() { | ||
final List<AirbyteStateType> allIncompatibleTypes = Lists.newArrayList(GLOBAL, STREAM); | ||
|
||
stateAggregator.ingest(getEmptyMessage(null)); | ||
|
||
allIncompatibleTypes.forEach(differentType -> Assertions.assertThatThrownBy(() -> stateAggregator.ingest(getEmptyMessage(differentType)))); | ||
|
||
stateAggregator.ingest(getEmptyMessage(LEGACY)); | ||
} | ||
|
||
@Test | ||
public void testNullState() { | ||
final AirbyteStateMessage state1 = getNullMessage(1); | ||
final AirbyteStateMessage state2 = getNullMessage(2); | ||
|
||
stateAggregator.ingest(state1); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(state1.getData())); | ||
|
||
stateAggregator.ingest(state2); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(state2.getData())); | ||
} | ||
|
||
@Test | ||
public void testLegacyState() { | ||
final AirbyteStateMessage state1 = getLegacyMessage(1); | ||
final AirbyteStateMessage state2 = getLegacyMessage(2); | ||
|
||
stateAggregator.ingest(state1); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(state1.getData())); | ||
|
||
stateAggregator.ingest(state2); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(state2.getData())); | ||
} | ||
|
||
@Test | ||
public void testGlobalState() { | ||
final AirbyteStateMessage state1 = getGlobalMessage(1); | ||
final AirbyteStateMessage state2 = getGlobalMessage(2); | ||
|
||
stateAggregator.ingest(state1); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(Jsons.jsonNode(List.of(state1)))); | ||
|
||
stateAggregator.ingest(state2); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(Jsons.jsonNode(List.of(state2)))); | ||
} | ||
|
||
@Test | ||
public void testStreamState() { | ||
final AirbyteStateMessage state1 = getStreamMessage("a", 1); | ||
final AirbyteStateMessage state2 = getStreamMessage("b", 2); | ||
final AirbyteStateMessage state3 = getStreamMessage("b", 3); | ||
|
||
stateAggregator.ingest(state1); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(Jsons.jsonNode(List.of(state1)))); | ||
|
||
stateAggregator.ingest(state2); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(Jsons.jsonNode(List.of(state2, state1)))); | ||
|
||
stateAggregator.ingest(state3); | ||
Assertions.assertThat(stateAggregator.getAggregated()).isEqualTo(new State() | ||
.withState(Jsons.jsonNode(List.of(state3, state1)))); | ||
} | ||
|
||
private AirbyteStateMessage getNullMessage(final int stateValue) { | ||
return new AirbyteStateMessage().withData(Jsons.jsonNode(stateValue)); | ||
} | ||
|
||
private AirbyteStateMessage getLegacyMessage(final int stateValue) { | ||
return new AirbyteStateMessage().withType(LEGACY).withData(Jsons.jsonNode(stateValue)); | ||
} | ||
|
||
private AirbyteStateMessage getGlobalMessage(final int stateValue) { | ||
return new AirbyteStateMessage().withType(GLOBAL) | ||
.withGlobal(new AirbyteGlobalState() | ||
.withStreamStates( | ||
List.of( | ||
new AirbyteStreamState() | ||
.withStreamDescriptor( | ||
new StreamDescriptor() | ||
.withName("test")) | ||
.withStreamState(Jsons.jsonNode(stateValue))))); | ||
} | ||
|
||
private AirbyteStateMessage getStreamMessage(final String streamName, final int stateValue) { | ||
return new AirbyteStateMessage().withType(STREAM) | ||
.withStream( | ||
new AirbyteStreamState() | ||
.withStreamDescriptor( | ||
new StreamDescriptor() | ||
.withName(streamName)) | ||
.withStreamState(Jsons.jsonNode(stateValue))); | ||
} | ||
|
||
private AirbyteStateMessage getEmptyMessage(final AirbyteStateType stateType) { | ||
if (stateType == STREAM) { | ||
return new AirbyteStateMessage() | ||
.withType(STREAM) | ||
.withStream( | ||
new AirbyteStreamState() | ||
.withStreamDescriptor(new StreamDescriptor())); | ||
} | ||
|
||
return new AirbyteStateMessage().withType(stateType); | ||
} | ||
|
||
} |