feat(Platform): update actor configuration when receiving control messages from connectors during sync #19811
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
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 |
|---|---|---|
|
|
@@ -26,6 +26,7 @@ | |
| import io.airbyte.config.WorkerDestinationConfig; | ||
| import io.airbyte.config.WorkerSourceConfig; | ||
| import io.airbyte.metrics.lib.ApmTraceUtils; | ||
| import io.airbyte.protocol.models.AirbyteControlMessage; | ||
| import io.airbyte.protocol.models.AirbyteMessage; | ||
| import io.airbyte.protocol.models.AirbyteMessage.Type; | ||
| import io.airbyte.protocol.models.AirbyteRecordMessage; | ||
|
|
@@ -36,6 +37,7 @@ | |
| import io.airbyte.workers.WorkerUtils; | ||
| import io.airbyte.workers.exception.RecordSchemaValidationException; | ||
| import io.airbyte.workers.exception.WorkerException; | ||
| import io.airbyte.workers.helper.ConnectorConfigUpdater; | ||
| import io.airbyte.workers.helper.FailureHelper; | ||
| import io.airbyte.workers.helper.ThreadedTimeTracker; | ||
| import io.airbyte.workers.internal.AirbyteDestination; | ||
|
|
@@ -98,6 +100,7 @@ public class DefaultReplicationWorker implements ReplicationWorker { | |
| private final AtomicBoolean hasFailed; | ||
| private final RecordSchemaValidator recordSchemaValidator; | ||
| private final WorkerMetricReporter metricReporter; | ||
| private final ConnectorConfigUpdater connectorConfigUpdater; | ||
| private final boolean fieldSelectionEnabled; | ||
|
|
||
| public DefaultReplicationWorker(final String jobId, | ||
|
|
@@ -108,6 +111,7 @@ public DefaultReplicationWorker(final String jobId, | |
| final MessageTracker messageTracker, | ||
| final RecordSchemaValidator recordSchemaValidator, | ||
| final WorkerMetricReporter metricReporter, | ||
| final ConnectorConfigUpdater connectorConfigUpdater, | ||
| final boolean fieldSelectionEnabled) { | ||
| this.jobId = jobId; | ||
| this.attempt = attempt; | ||
|
|
@@ -118,6 +122,7 @@ public DefaultReplicationWorker(final String jobId, | |
| this.executors = Executors.newFixedThreadPool(2); | ||
| this.recordSchemaValidator = recordSchemaValidator; | ||
| this.metricReporter = metricReporter; | ||
| this.connectorConfigUpdater = connectorConfigUpdater; | ||
| this.fieldSelectionEnabled = fieldSelectionEnabled; | ||
|
|
||
| this.cancelled = new AtomicBoolean(false); | ||
|
|
@@ -191,7 +196,7 @@ private void replicate(final Path jobRoot, | |
| // note: `whenComplete` is used instead of `exceptionally` so that the original exception is still | ||
| // thrown | ||
| final CompletableFuture<?> readFromDstThread = CompletableFuture.runAsync( | ||
| readFromDstRunnable(destination, cancelled, messageTracker, connectorConfigUpdater, mdc, timeTracker, destinationConfig), | ||
| executors) | ||
| .whenComplete((msg, ex) -> { | ||
| if (ex != null) { | ||
|
|
@@ -212,10 +217,12 @@ private void replicate(final Path jobRoot, | |
| cancelled, | ||
| mapper, | ||
| messageTracker, | ||
| connectorConfigUpdater, | ||
| mdc, | ||
| recordSchemaValidator, | ||
| metricReporter, | ||
| timeTracker, | ||
| sourceConfig, | ||
| fieldSelectionEnabled), | ||
| executors) | ||
| .whenComplete((msg, ex) -> { | ||
|
|
@@ -253,8 +260,10 @@ private void replicate(final Path jobRoot, | |
| private static Runnable readFromDstRunnable(final AirbyteDestination destination, | ||
| final AtomicBoolean cancelled, | ||
| final MessageTracker messageTracker, | ||
| final ConnectorConfigUpdater connectorConfigUpdater, | ||
| final Map<String, String> mdc, | ||
| final ThreadedTimeTracker timeHolder, | ||
| final WorkerDestinationConfig destinationConfig) { | ||
| return () -> { | ||
| MDC.setContextMap(mdc); | ||
| LOGGER.info("Destination output thread started."); | ||
|
|
@@ -267,8 +276,18 @@ private static Runnable readFromDstRunnable(final AirbyteDestination destination | |
| throw new DestinationException("Destination process read attempt failed", e); | ||
| } | ||
| if (messageOptional.isPresent()) { | ||
| final AirbyteMessage message = messageOptional.get(); | ||
| LOGGER.info("State in DefaultReplicationWorker from destination: {}", message); | ||
|
|
||
| messageTracker.acceptFromDestination(message); | ||
|
|
||
| try { | ||
| if (message.getType() == Type.CONTROL) { | ||
| acceptDstControlMessage(destinationConfig, message.getControl(), connectorConfigUpdater); | ||
| } | ||
| } catch (final Exception e) { | ||
| LOGGER.error("Error updating destination configuration", e); | ||
| } | ||
| } | ||
| } | ||
| timeHolder.trackDestinationWriteEndTime(); | ||
|
|
@@ -300,10 +319,12 @@ private static Runnable readFromSrcAndWriteToDstRunnable(final AirbyteSource sou | |
| final AtomicBoolean cancelled, | ||
| final AirbyteMapper mapper, | ||
| final MessageTracker messageTracker, | ||
| final ConnectorConfigUpdater connectorConfigUpdater, | ||
| final Map<String, String> mdc, | ||
| final RecordSchemaValidator recordSchemaValidator, | ||
| final WorkerMetricReporter metricReporter, | ||
| final ThreadedTimeTracker timeHolder, | ||
| final WorkerSourceConfig sourceConfig, | ||
| final boolean fieldSelectionEnabled) { | ||
| return () -> { | ||
| MDC.setContextMap(mdc); | ||
|
|
@@ -333,6 +354,14 @@ private static Runnable readFromSrcAndWriteToDstRunnable(final AirbyteSource sou | |
|
|
||
| messageTracker.acceptFromSource(message); | ||
|
|
||
| try { | ||
| if (message.getType() == Type.CONTROL) { | ||
| acceptSrcControlMessage(sourceConfig, message.getControl(), connectorConfigUpdater); | ||
| } | ||
| } catch (final Exception e) { | ||
| LOGGER.error("Error updating source configuration", e); | ||
| } | ||
|
|
||
| try { | ||
| if (message.getType() == Type.RECORD || message.getType() == Type.STATE) { | ||
| destination.accept(message); | ||
|
|
@@ -391,6 +420,22 @@ private static Runnable readFromSrcAndWriteToDstRunnable(final AirbyteSource sou | |
| }; | ||
| } | ||
|
|
||
| private static void acceptSrcControlMessage(final WorkerSourceConfig sourceConfig, | ||
|
There was a problem hiding this comment. we can simplify this method signature by passing in the id instead of the entire config There was a problem hiding this comment. this simplification also extends to the various runnable method interfaces e.g There was a problem hiding this comment. 👍🏻 updated in d6eb05f |
||
| final AirbyteControlMessage controlMessage, | ||
| final ConnectorConfigUpdater connectorConfigUpdater) { | ||
| if (controlMessage.getType() == AirbyteControlMessage.Type.CONNECTOR_CONFIG) { | ||
|
There was a problem hiding this comment. are there other kinds of control messages today? There was a problem hiding this comment. Nope, this is the only one for now! |
||
| connectorConfigUpdater.updateSource(sourceConfig.getSourceId(), controlMessage.getConnectorConfig().getConfig()); | ||
| } | ||
| } | ||
|
|
||
| private static void acceptDstControlMessage(final WorkerDestinationConfig destinationConfig, | ||
| final AirbyteControlMessage controlMessage, | ||
| final ConnectorConfigUpdater connectorConfigUpdater) { | ||
| if (controlMessage.getType() == AirbyteControlMessage.Type.CONNECTOR_CONFIG) { | ||
| connectorConfigUpdater.updateDestination(destinationConfig.getDestinationId(), controlMessage.getConnectorConfig().getConfig()); | ||
| } | ||
| } | ||
|
|
||
| private ReplicationOutput getReplicationOutput(final StandardSyncInput syncInput, | ||
| final WorkerDestinationConfig destinationConfig, | ||
| final AtomicReference<FailureReason> replicationRunnableFailureRef, | ||
|
|
||
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,87 @@ | ||
| /* | ||
| * Copyright (c) 2022 Airbyte, Inc., all rights reserved. | ||
| */ | ||
|
|
||
| package io.airbyte.workers.helper; | ||
|
|
||
| import com.google.common.hash.Hashing; | ||
| import io.airbyte.api.client.AirbyteApiClient; | ||
| import io.airbyte.api.client.generated.DestinationApi; | ||
| import io.airbyte.api.client.generated.SourceApi; | ||
| import io.airbyte.api.client.model.generated.DestinationIdRequestBody; | ||
| import io.airbyte.api.client.model.generated.DestinationRead; | ||
| import io.airbyte.api.client.model.generated.DestinationUpdate; | ||
| import io.airbyte.api.client.model.generated.SourceIdRequestBody; | ||
| import io.airbyte.api.client.model.generated.SourceRead; | ||
| import io.airbyte.api.client.model.generated.SourceUpdate; | ||
| import io.airbyte.commons.json.Jsons; | ||
| import io.airbyte.protocol.models.Config; | ||
| import java.nio.charset.StandardCharsets; | ||
| import java.util.UUID; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| /** | ||
|
|
||
| * Helper class for workers to persist updates to Source/Destination configs emitted from | ||
| * AirbyteControlMessages. | ||
| * | ||
| * This is in order to support connectors updating configs when running commands, which is specially | ||
| * useful for migrating configuration to a new version or for enabling connectors that require | ||
| * single-use or short-lived OAuth tokens. | ||
| */ | ||
| public class ConnectorConfigUpdater { | ||
|
|
||
| private static final Logger LOGGER = LoggerFactory.getLogger(ConnectorConfigUpdater.class); | ||
|
|
||
| private final SourceApi sourceApi; | ||
| private final DestinationApi destinationApi; | ||
|
|
||
| public ConnectorConfigUpdater(final SourceApi sourceApi, final DestinationApi destinationApi) { | ||
| this.sourceApi = sourceApi; | ||
| this.destinationApi = destinationApi; | ||
| } | ||
|
|
||
| /** | ||
| * Updates the Source from a sync job ID with the provided Configuration. Secrets and OAuth | ||
| * parameters will be masked when saving. | ||
| */ | ||
| public void updateSource(final UUID sourceId, final Config config) { | ||
| final SourceRead source = AirbyteApiClient.retryWithJitter( | ||
| () -> sourceApi.getSource(new SourceIdRequestBody().sourceId(sourceId)), | ||
| "get source"); | ||
|
|
||
| final SourceRead updatedSource = AirbyteApiClient.retryWithJitter( | ||
| () -> sourceApi | ||
| .updateSource(new SourceUpdate() | ||
| .sourceId(sourceId) | ||
| .name(source.getName()) | ||
| .connectionConfiguration(Jsons.jsonNode(config.getAdditionalProperties()))), | ||
| "update source"); | ||
|
|
||
| LOGGER.info("Persisted updated configuration for source {}. New config hash: {}.", sourceId, | ||
| Hashing.sha256().hashString(updatedSource.getConnectionConfiguration().asText(), StandardCharsets.UTF_8)); | ||
|
|
||
| } | ||
|
|
||
| /** | ||
| * Updates the Destination from a sync job ID with the provided Configuration. Secrets and OAuth | ||
| * parameters will be masked when saving. | ||
| */ | ||
| public void updateDestination(final UUID destinationId, final Config config) { | ||
| final DestinationRead destination = AirbyteApiClient.retryWithJitter( | ||
| () -> destinationApi.getDestination(new DestinationIdRequestBody().destinationId(destinationId)), | ||
| "get destination"); | ||
|
|
||
| final DestinationRead updatedDestination = AirbyteApiClient.retryWithJitter( | ||
| () -> destinationApi | ||
| .updateDestination(new DestinationUpdate() | ||
| .destinationId(destinationId) | ||
| .name(destination.getName()) | ||
| .connectionConfiguration(Jsons.jsonNode(config.getAdditionalProperties()))), | ||
| "update destination"); | ||
|
|
||
| LOGGER.info("Persisted updated configuration for destination {}. New config hash: {}.", destinationId, | ||
| Hashing.sha256().hashString(updatedDestination.getConnectionConfiguration().asText(), StandardCharsets.UTF_8)); | ||
| } | ||
|
|
||
| } | ||
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
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
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.
curious, what is this for?
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.
This was necessary in order to move the ApiClientBeanFactory into commons-worker (to reuse it in both airbyte-workers and container-orchestrator). It does some token generation that needs this dependency.
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.
This might be an issue, we split up
airbyte-commons-workerrecently due to some initialization issues of the beans when we added cron.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.
@gosusnp Could you expand a bit on what the issue might be? My change moved the
AirbyteAPiClientBeanFactoryfromairbyte-workersintoairbyte-commons-worker. I don't seeairbyte-workersnorairbyte-commons-workerbeing a dependency ofairbyte-cron.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.
I was looking for an issue to link for better context but I cannot find it...
I think the problem we had is that micronaut discovers all singleton upon starts, so by having some services in a shared lib led us to start more services than intended. For example, the cron would pick up some beans from the workers and would effectively behave as an extra worker.
@jdpgrailsdev or @benmoriceau might have better context since they were more closely working on this.