-
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 destinations to handle new state messages (#13670)
- Loading branch information
Showing
8 changed files
with
751 additions
and
20 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
114 changes: 114 additions & 0 deletions
114
...tegrations/destination/dest_state_lifecycle_manager/DefaultDestStateLifecycleManager.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,114 @@ | ||
| /* | ||
| * Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
| */ | ||
|
|
||
| package io.airbyte.integrations.destination.dest_state_lifecycle_manager; | ||
|
|
||
| import com.google.common.annotations.VisibleForTesting; | ||
| import com.google.common.base.Preconditions; | ||
| import io.airbyte.protocol.models.AirbyteMessage; | ||
| import io.airbyte.protocol.models.AirbyteMessage.Type; | ||
| import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType; | ||
| import java.util.Queue; | ||
| import java.util.function.Supplier; | ||
|
|
||
| /** | ||
| * Detects the type of the state being received by anchoring on the first state type it sees. Fail | ||
| * if receives states of multiple types--each instance of this class can only support state messages | ||
| * of one type. The protocol specifies that a source should emit state messages of a single type | ||
| * during a sync, so a single instance of this manager is sufficient for a destination to track | ||
| * state during a sync. | ||
| * | ||
| * Strategy: Delegates state messages of each type to a StateManager that is appropriate to that | ||
| * state type. | ||
| * | ||
| * Per the protocol, if state type is not set, assumes the LEGACY state type. | ||
| */ | ||
| public class DefaultDestStateLifecycleManager implements DestStateLifecycleManager { | ||
|
|
||
| private AirbyteStateType stateType; | ||
| private final Supplier<DestStateLifecycleManager> internalStateManagerSupplier; | ||
|
|
||
| public DefaultDestStateLifecycleManager() { | ||
| this(new DestSingleStateLifecycleManager(), new DestStreamStateLifecycleManager()); | ||
| } | ||
|
|
||
| @VisibleForTesting | ||
| DefaultDestStateLifecycleManager(final DestStateLifecycleManager singleStateManager, final DestStateLifecycleManager streamStateManager) { | ||
| stateType = null; | ||
| // allows us to delegate calls to the appropriate underlying state manager. | ||
| internalStateManagerSupplier = () -> { | ||
| if (stateType == AirbyteStateType.GLOBAL || stateType == AirbyteStateType.LEGACY || stateType == null) { | ||
| return singleStateManager; | ||
| } else if (stateType == AirbyteStateType.STREAM) { | ||
| return streamStateManager; | ||
| } else { | ||
| throw new IllegalArgumentException("unrecognized state type"); | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| @Override | ||
| public void addState(final AirbyteMessage message) { | ||
| Preconditions.checkArgument(message.getType() == Type.STATE, "Messages passed to State Manager must be of type STATE."); | ||
| Preconditions.checkArgument(isStateTypeCompatible(stateType, message.getState().getStateType())); | ||
|
|
||
| setManagerStateTypeIfNotSet(message); | ||
|
|
||
| internalStateManagerSupplier.get().addState(message); | ||
| } | ||
|
|
||
| /** | ||
| * Given the type of previously recorded state by the state manager, determines if a newly added | ||
| * state message's type is compatible. Based on the previously set state type, determines if a new | ||
| * one is compatible. If the previous state is null, any new state is compatible. If new state type | ||
| * is null, it should be treated as LEGACY. Thus, previousStateType == LEGACY and newStateType == | ||
| * null IS compatible. All other state types are compatible based on equality. | ||
| * | ||
| * @param previousStateType - state type previously recorded by the state manager | ||
| * @param newStateType - state message of a newly added message | ||
| * @return true if compatible, otherwise false | ||
| */ | ||
| private static boolean isStateTypeCompatible(final AirbyteStateType previousStateType, final AirbyteStateType newStateType) { | ||
| return previousStateType == null || previousStateType == AirbyteStateType.LEGACY && newStateType == null || previousStateType == newStateType; | ||
| } | ||
|
|
||
| /** | ||
| * If the state type for the manager is not set, sets it using the state type from the message. If | ||
| * the type on the message is null, we assume it is LEGACY. After the first, state message is added | ||
| * to the manager, the state type is set and is immutable. | ||
| * | ||
| * @param message - state message whose state will be used if internal state type is not set | ||
| */ | ||
| private void setManagerStateTypeIfNotSet(final AirbyteMessage message) { | ||
| // detect and set state type. | ||
| if (stateType == null) { | ||
| if (message.getState().getStateType() == null) { | ||
| stateType = AirbyteStateType.LEGACY; | ||
| } else { | ||
| stateType = message.getState().getStateType(); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void markPendingAsFlushed() { | ||
| internalStateManagerSupplier.get().markPendingAsFlushed(); | ||
| } | ||
|
|
||
| @Override | ||
| public Queue<AirbyteMessage> listFlushed() { | ||
| return internalStateManagerSupplier.get().listFlushed(); | ||
| } | ||
|
|
||
| @Override | ||
| public void markFlushedAsCommitted() { | ||
| internalStateManagerSupplier.get().markFlushedAsCommitted(); | ||
| } | ||
|
|
||
| @Override | ||
| public Queue<AirbyteMessage> listCommitted() { | ||
| return internalStateManagerSupplier.get().listCommitted(); | ||
| } | ||
|
|
||
| } |
68 changes: 68 additions & 0 deletions
68
...ntegrations/destination/dest_state_lifecycle_manager/DestSingleStateLifecycleManager.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,68 @@ | ||
| /* | ||
| * Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
| */ | ||
|
|
||
| package io.airbyte.integrations.destination.dest_state_lifecycle_manager; | ||
|
|
||
| import com.google.common.annotations.VisibleForTesting; | ||
| import io.airbyte.protocol.models.AirbyteMessage; | ||
| import java.util.Collections; | ||
| import java.util.LinkedList; | ||
| import java.util.List; | ||
| import java.util.Queue; | ||
|
|
||
| /** | ||
| * This {@link DestStateLifecycleManager} handles any state where there is a guarantee that any | ||
| * single state message represents the state for the ENTIRE connection. At the time of writing, | ||
| * GLOBAL and LEGACY state types are the state type that match this pattern. | ||
| * | ||
| * Does NOT store duplicates. Because each state message represents the entire state for the | ||
| * connection, it only stores (and emits) the LAST state it received at each phase. | ||
| */ | ||
| public class DestSingleStateLifecycleManager implements DestStateLifecycleManager { | ||
|
|
||
| private AirbyteMessage lastPendingState; | ||
| private AirbyteMessage lastFlushedState; | ||
| private AirbyteMessage lastCommittedState; | ||
|
|
||
| @Override | ||
| public void addState(final AirbyteMessage message) { | ||
| lastPendingState = message; | ||
| } | ||
|
|
||
| @VisibleForTesting | ||
| Queue<AirbyteMessage> listPending() { | ||
| return stateMessageToQueue(lastPendingState); | ||
| } | ||
|
|
||
| @Override | ||
| public void markPendingAsFlushed() { | ||
| if (lastPendingState != null) { | ||
| lastFlushedState = lastPendingState; | ||
| lastPendingState = null; | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public Queue<AirbyteMessage> listFlushed() { | ||
| return stateMessageToQueue(lastFlushedState); | ||
| } | ||
|
|
||
| @Override | ||
| public void markFlushedAsCommitted() { | ||
| if (lastFlushedState != null) { | ||
| lastCommittedState = lastFlushedState; | ||
| lastFlushedState = null; | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public Queue<AirbyteMessage> listCommitted() { | ||
| return stateMessageToQueue(lastCommittedState); | ||
| } | ||
|
|
||
| private static Queue<AirbyteMessage> stateMessageToQueue(final AirbyteMessage stateMessage) { | ||
| return new LinkedList<>(stateMessage == null ? Collections.emptyList() : List.of(stateMessage)); | ||
| } | ||
|
|
||
| } |
53 changes: 53 additions & 0 deletions
53
...byte/integrations/destination/dest_state_lifecycle_manager/DestStateLifecycleManager.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,53 @@ | ||
| /* | ||
| * Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
| */ | ||
|
|
||
| package io.airbyte.integrations.destination.dest_state_lifecycle_manager; | ||
|
|
||
| import io.airbyte.protocol.models.AirbyteMessage; | ||
| import java.util.Queue; | ||
|
|
||
| /** | ||
| * This class manages the lifecycle of state message. It tracks state messages that are in 3 states: | ||
| * <ol> | ||
| * <li>pending - associated records have been accepted by the connector but has NOT been pushed to | ||
| * the destination</li> | ||
| * <li>flushed - associated records have been flushed to tmp storage in the destination but have NOT | ||
| * been committed</li> | ||
| * <li>committed - associated records have been committed</li> | ||
| * </ol> | ||
| */ | ||
| public interface DestStateLifecycleManager { | ||
|
|
||
| /** | ||
| * Accepts a state into the manager. The state starts in a pending state. | ||
| * | ||
| * @param message - airbyte message of type state | ||
| */ | ||
| void addState(AirbyteMessage message); | ||
|
|
||
| /** | ||
| * Moves any tracked state messages that are currently pending to flushed. | ||
| */ | ||
| void markPendingAsFlushed(); | ||
|
|
||
| /** | ||
| * List all tracked state messages that are flushed. | ||
| * | ||
| * @return list of state messages | ||
| */ | ||
| Queue<AirbyteMessage> listFlushed(); | ||
|
|
||
| /** | ||
| * Moves any tracked state messages that are currently flushed to committed. | ||
| */ | ||
| void markFlushedAsCommitted(); | ||
|
|
||
| /** | ||
| * List all tracked state messages that are committed. | ||
| * | ||
| * @return list of state messages | ||
| */ | ||
| Queue<AirbyteMessage> listCommitted(); | ||
|
|
||
| } |
Oops, something went wrong.