process config control messages during check and discover

airbyte-api/src/main/openapi/config.yaml

@@ -2667,6 +2667,8 @@ components:
         - connectionConfiguration
         - workspaceId
       properties:
+        sourceId:
+          $ref: "#/components/schemas/SourceId"
         sourceDefinitionId:
           $ref: "#/components/schemas/SourceDefinitionId"
         connectionConfiguration:
@@ -2959,6 +2961,8 @@ components:
         - destinationDefinitionId
         - connectionConfiguration
       properties:
+        destinationId:
+          $ref: "#/components/schemas/DestinationId"
         destinationDefinitionId:
           $ref: "#/components/schemas/DestinationDefinitionId"
         connectionConfiguration:

airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java

@@ -349,7 +349,10 @@ public TemporalResponse<ConnectorJobOutput> submitCheckConnection(final UUID job
         .withDockerImage(config.getDockerImage())
         .withProtocolVersion(config.getProtocolVersion())
         .withIsCustomConnector(config.getIsCustomConnector());
-    final StandardCheckConnectionInput input = new StandardCheckConnectionInput().withConnectionConfiguration(config.getConnectionConfiguration());
+    final StandardCheckConnectionInput input = new StandardCheckConnectionInput()
+        .withActorType(config.getActorType())
+        .withActorId(config.getActorId())
+        .withConnectionConfiguration(config.getConnectionConfiguration());
 
     return execute(jobRunConfig,
         () -> getWorkflowStub(CheckConnectionWorkflow.class, TemporalJobType.CHECK_CONNECTION).run(jobRunConfig, launcherConfig, input));

airbyte-commons-worker/build.gradle

@@ -10,6 +10,7 @@ dependencies {
     implementation libs.bundles.micronaut
 
     implementation 'io.fabric8:kubernetes-client:5.12.2'
+    implementation 'com.auth0:java-jwt:3.19.2'
     implementation libs.guava
     implementation (libs.temporal.sdk) {
         exclude module: 'guava'

airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerUtils.java

@@ -11,6 +11,8 @@
 import io.airbyte.config.StandardSyncInput;
 import io.airbyte.config.WorkerDestinationConfig;
 import io.airbyte.config.WorkerSourceConfig;
+import io.airbyte.protocol.models.AirbyteControlConnectorConfigMessage;
+import io.airbyte.protocol.models.AirbyteControlMessage;
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteMessage.Type;
 import io.airbyte.protocol.models.AirbyteStreamNameNamespacePair;
@@ -91,6 +93,7 @@ public static void cancelProcess(final Process process) {
    */
   public static WorkerSourceConfig syncToWorkerSourceConfig(final StandardSyncInput sync) {
     return new WorkerSourceConfig()
+        .withSourceId(sync.getSourceId())
         .withSourceConnectionConfiguration(sync.getSourceConfiguration())
         .withCatalog(sync.getCatalog())
         .withState(sync.getState());
@@ -102,11 +105,22 @@ public static WorkerSourceConfig syncToWorkerSourceConfig(final StandardSyncInpu
    */
   public static WorkerDestinationConfig syncToWorkerDestinationConfig(final StandardSyncInput sync) {
     return new WorkerDestinationConfig()
+        .withDestinationId(sync.getDestinationId())
         .withDestinationConnectionConfiguration(sync.getDestinationConfiguration())
         .withCatalog(sync.getCatalog())
         .withState(sync.getState());
   }
 
+  public static Optional<AirbyteControlConnectorConfigMessage> getMostRecentConfigControlMessage(final Map<Type, List<AirbyteMessage>> messagesByType) {
+    final List<AirbyteControlConnectorConfigMessage> configMessages = messagesByType.getOrDefault(Type.CONTROL, new ArrayList<>()).stream()
+        .map(AirbyteMessage::getControl)
+        .filter(control -> control.getType() == AirbyteControlMessage.Type.CONNECTOR_CONFIG)
+        .map(AirbyteControlMessage::getConnectorConfig)
+        .toList();
+
+    return configMessages.isEmpty() ? Optional.empty() : Optional.of(configMessages.get(configMessages.size() - 1));
+  }
+
   public static ConnectorJobOutput getJobFailureOutputOrThrow(final OutputType outputType,
                                                               final Map<Type, List<AirbyteMessage>> messagesByType,
                                                               final String defaultErrorMessage)

airbyte-workers/src/main/java/io/airbyte/workers/config/ApiClientBeanFactory.java → airbyte-commons-worker/src/main/java/io/airbyte/workers/config/ApiClientBeanFactory.java

@@ -9,6 +9,7 @@
 import com.auth0.jwt.algorithms.Algorithm;
 import com.google.auth.oauth2.ServiceAccountCredentials;
 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.invoker.generated.ApiClient;
 import io.airbyte.commons.temporal.config.WorkerMode;
@@ -58,7 +59,7 @@ public ApiClient apiClient(@Value("${airbyte.internal.api.auth-header.name}") fi
   }
 
   @Singleton
-  public AirbyteApiClient airbyteApiClient(ApiClient apiClient) {
+  public AirbyteApiClient airbyteApiClient(final ApiClient apiClient) {
     return new AirbyteApiClient(apiClient);
   }
 
@@ -67,6 +68,11 @@ public SourceApi sourceApi(final ApiClient apiClient) {
     return new SourceApi(apiClient);
   }
 
+  @Singleton
+  public DestinationApi destinationApi(final ApiClient apiClient) {
+    return new DestinationApi(apiClient);
+  }
+
   @Singleton
   public HttpClient httpClient() {
     return HttpClient.newHttpClient();

airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DefaultCheckConnectionWorker.java

@@ -19,11 +19,13 @@
 import io.airbyte.config.StandardCheckConnectionOutput.Status;
 import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.protocol.models.AirbyteConnectionStatus;
+import io.airbyte.protocol.models.AirbyteControlConnectorConfigMessage;
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteMessage.Type;
 import io.airbyte.workers.WorkerConstants;
 import io.airbyte.workers.WorkerUtils;
 import io.airbyte.workers.exception.WorkerException;
+import io.airbyte.workers.helper.ConnectorConfigUpdater;
 import io.airbyte.workers.internal.AirbyteStreamFactory;
 import io.airbyte.workers.internal.DefaultAirbyteStreamFactory;
 import io.airbyte.workers.process.IntegrationLauncher;
@@ -43,18 +45,21 @@ public class DefaultCheckConnectionWorker implements CheckConnectionWorker {
   private static final Logger LOGGER = LoggerFactory.getLogger(DefaultCheckConnectionWorker.class);
 
   private final IntegrationLauncher integrationLauncher;
+  private final ConnectorConfigUpdater connectorConfigUpdater;
   private final AirbyteStreamFactory streamFactory;
 
   private Process process;
 
   public DefaultCheckConnectionWorker(final IntegrationLauncher integrationLauncher,
+                                      final ConnectorConfigUpdater connectorConfigUpdater,
                                       final AirbyteStreamFactory streamFactory) {
     this.integrationLauncher = integrationLauncher;
+    this.connectorConfigUpdater = connectorConfigUpdater;
     this.streamFactory = streamFactory;
   }
 
-  public DefaultCheckConnectionWorker(final IntegrationLauncher integrationLauncher) {
-    this(integrationLauncher, new DefaultAirbyteStreamFactory());
+  public DefaultCheckConnectionWorker(final IntegrationLauncher integrationLauncher, final ConnectorConfigUpdater connectorConfigUpdater) {
+    this(integrationLauncher, connectorConfigUpdater, new DefaultAirbyteStreamFactory());
   }
 
   @Trace(operationName = WORKER_OPERATION_NAME)
@@ -84,6 +89,21 @@ public ConnectorJobOutput run(final StandardCheckConnectionInput input, final Pa
           .map(AirbyteMessage::getConnectionStatus)
           .findFirst();
 
+      if (input.getActorId() != null && input.getActorType() != null) {
+        final Optional<AirbyteControlConnectorConfigMessage> optionalConfigMsg = WorkerUtils.getMostRecentConfigControlMessage(messagesByType);
+        optionalConfigMsg.ifPresent(
+            configMessage -> {
+              switch (input.getActorType()) {
+                case SOURCE -> connectorConfigUpdater.updateSource(
+                    input.getActorId(),
+                    configMessage.getConfig());
+                case DESTINATION -> connectorConfigUpdater.updateDestination(
+                    input.getActorId(),
+                    configMessage.getConfig());
+              }
+            });
+      }
+
       if (status.isPresent() && exitCode == 0) {
         final StandardCheckConnectionOutput output = new StandardCheckConnectionOutput()
             .withStatus(Enums.convertTo(status.get().getStatus(), Status.class))

airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DefaultDiscoverCatalogWorker.java

@@ -19,11 +19,13 @@
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.protocol.models.AirbyteCatalog;
+import io.airbyte.protocol.models.AirbyteControlConnectorConfigMessage;
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteMessage.Type;
 import io.airbyte.workers.WorkerConstants;
 import io.airbyte.workers.WorkerUtils;
 import io.airbyte.workers.exception.WorkerException;
+import io.airbyte.workers.helper.ConnectorConfigUpdater;
 import io.airbyte.workers.internal.AirbyteStreamFactory;
 import io.airbyte.workers.internal.DefaultAirbyteStreamFactory;
 import io.airbyte.workers.process.IntegrationLauncher;
@@ -48,20 +50,24 @@ public class DefaultDiscoverCatalogWorker implements DiscoverCatalogWorker {
 
   private final IntegrationLauncher integrationLauncher;
   private final AirbyteStreamFactory streamFactory;
+  private final ConnectorConfigUpdater connectorConfigUpdater;
 
   private volatile Process process;
 
   public DefaultDiscoverCatalogWorker(final ConfigRepository configRepository,
                                       final IntegrationLauncher integrationLauncher,
+                                      final ConnectorConfigUpdater connectorConfigUpdater,
                                       final AirbyteStreamFactory streamFactory) {
     this.configRepository = configRepository;
     this.integrationLauncher = integrationLauncher;
     this.streamFactory = streamFactory;
+    this.connectorConfigUpdater = connectorConfigUpdater;
   }
 
   public DefaultDiscoverCatalogWorker(final ConfigRepository configRepository,
-                                      final IntegrationLauncher integrationLauncher) {
-    this(configRepository, integrationLauncher, new DefaultAirbyteStreamFactory());
+                                      final IntegrationLauncher integrationLauncher,
+                                      final ConnectorConfigUpdater connectorConfigUpdater) {
+    this(configRepository, integrationLauncher, connectorConfigUpdater, new DefaultAirbyteStreamFactory());
   }
 
   @Trace(operationName = WORKER_OPERATION_NAME)
@@ -90,6 +96,12 @@ public ConnectorJobOutput run(final StandardDiscoverCatalogInput discoverSchemaI
           .map(AirbyteMessage::getCatalog)
           .findFirst();
 
+      final Optional<AirbyteControlConnectorConfigMessage> optionalConfigMsg = WorkerUtils.getMostRecentConfigControlMessage(messagesByType);
+      optionalConfigMsg.ifPresent(
+          configMessage -> connectorConfigUpdater.updateSource(
+              UUID.fromString(discoverSchemaInput.getSourceId()),
+              configMessage.getConfig()));
+
       final int exitCode = process.exitValue();
       if (exitCode == 0) {
         if (catalog.isEmpty()) {

airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DefaultReplicationWorker.java

@@ -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, mdc, timeTracker),
+          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 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()) {
-            LOGGER.info("State in DefaultReplicationWorker from destination: {}", messageOptional.get());
-            messageTracker.acceptFromDestination(messageOptional.get());
+            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,
+                                              final AirbyteControlMessage controlMessage,
+                                              final ConnectorConfigUpdater connectorConfigUpdater) {
+    if (controlMessage.getType() == AirbyteControlMessage.Type.CONNECTOR_CONFIG) {
+      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,