Re-introduce changes to add new labels to sync kubernetes pods (#7719)

Now that backwards compatibility is enabled in IntegrationLauncherConfig, we can safely add new properties to enable reporting workspace id and connection id on sync pods. Transforms are not yet supported.

airbyte-commons-server/src/main/java/io/airbyte/commons/server/handlers/JobInputHandler.java

@@ -185,13 +185,15 @@ public Object getJobInput(final SyncInput input) {
       final IntegrationLauncherConfig sourceLauncherConfig = getSourceIntegrationLauncherConfig(
           jobId,
           attempt,
+          connectionId,
           config,
           sourceVersion,
           attemptSyncConfig.getSourceConfiguration());
 
       final IntegrationLauncherConfig destinationLauncherConfig = getDestinationIntegrationLauncherConfig(
           jobId,
           attempt,
+          connectionId,
           config,
           destinationVersion,
           attemptSyncConfig.getDestinationConfiguration(),
@@ -218,7 +220,7 @@ public Object getJobInput(final SyncInput input) {
           .withResourceRequirements(config.getResourceRequirements())
           .withSourceResourceRequirements(config.getSourceResourceRequirements())
           .withDestinationResourceRequirements(config.getDestinationResourceRequirements())
-          .withConnectionId(standardSync.getConnectionId())
+          .withConnectionId(connectionId)
           .withWorkspaceId(config.getWorkspaceId())
           .withNormalizeInDestinationContainer(shouldNormalizeInDestination)
           .withIsReset(JobConfig.ConfigType.RESET_CONNECTION.equals(jobConfigType));
@@ -263,6 +265,7 @@ public Object getCheckJobInput(final CheckInput input) {
       final IntegrationLauncherConfig sourceLauncherConfig = getSourceIntegrationLauncherConfig(
           jobId,
           attemptNumber,
+          connectionId,
           jobSyncConfig,
           sourceVersion,
           sourceConfiguration);
@@ -271,6 +274,7 @@ public Object getCheckJobInput(final CheckInput input) {
           getDestinationIntegrationLauncherConfig(
               jobId,
               attemptNumber,
+              connectionId,
               jobSyncConfig,
               destinationVersion,
               destinationConfiguration,
@@ -361,6 +365,7 @@ private void reportNormalizationInDestinationMetrics(final boolean shouldNormali
 
   private IntegrationLauncherConfig getSourceIntegrationLauncherConfig(final long jobId,
                                                                        final int attempt,
+                                                                       final UUID connectionId,
                                                                        final JobSyncConfig config,
                                                                        @Nullable final ActorDefinitionVersion sourceVersion,
                                                                        final JsonNode sourceConfiguration)
@@ -370,6 +375,8 @@ private IntegrationLauncherConfig getSourceIntegrationLauncherConfig(final long
     final IntegrationLauncherConfig sourceLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(jobId))
         .withAttemptId((long) attempt)
+        .withConnectionId(connectionId)
+        .withWorkspaceId(config.getWorkspaceId())
         .withDockerImage(config.getSourceDockerImage())
         .withProtocolVersion(config.getSourceProtocolVersion())
         .withIsCustomConnector(config.getIsSourceCustomConnector());
@@ -383,6 +390,7 @@ private IntegrationLauncherConfig getSourceIntegrationLauncherConfig(final long
 
   private IntegrationLauncherConfig getDestinationIntegrationLauncherConfig(final long jobId,
                                                                             final int attempt,
+                                                                            final UUID connectionId,
                                                                             final JobSyncConfig config,
                                                                             final ActorDefinitionVersion destinationVersion,
                                                                             final JsonNode destinationConfiguration,
@@ -400,6 +408,8 @@ private IntegrationLauncherConfig getDestinationIntegrationLauncherConfig(final
     return new IntegrationLauncherConfig()
         .withJobId(String.valueOf(jobId))
         .withAttemptId((long) attempt)
+        .withConnectionId(connectionId)
+        .withWorkspaceId(config.getWorkspaceId())
         .withDockerImage(config.getDestinationDockerImage())
         .withProtocolVersion(config.getDestinationProtocolVersion())
         .withIsCustomConnector(config.getIsDestinationCustomConnector())

airbyte-commons-server/src/main/java/io/airbyte/commons/server/scheduler/DefaultSynchronousSchedulerClient.java

@@ -103,7 +103,7 @@ public SynchronousResponse<StandardCheckConnectionOutput> createSourceCheckConne
         ConfigType.CHECK_CONNECTION_SOURCE,
         jobReportingContext,
         source.getSourceDefinitionId(),
-        () -> temporalClient.submitCheckConnection(UUID.randomUUID(), 0, taskQueue, jobCheckConnectionConfig),
+        () -> temporalClient.submitCheckConnection(UUID.randomUUID(), 0, source.getWorkspaceId(), taskQueue, jobCheckConnectionConfig),
         ConnectorJobOutput::getCheckConnection,
         source.getWorkspaceId(),
         source.getSourceId());
@@ -137,7 +137,7 @@ public SynchronousResponse<StandardCheckConnectionOutput> createDestinationCheck
         ConfigType.CHECK_CONNECTION_DESTINATION,
         jobReportingContext,
         destination.getDestinationDefinitionId(),
-        () -> temporalClient.submitCheckConnection(jobId, 0, taskQueue, jobCheckConnectionConfig),
+        () -> temporalClient.submitCheckConnection(jobId, 0, destination.getWorkspaceId(), taskQueue, jobCheckConnectionConfig),
         ConnectorJobOutput::getCheckConnection,
         destination.getWorkspaceId(),
         destination.getDestinationId());
@@ -173,7 +173,7 @@ public SynchronousResponse<UUID> createDiscoverSchemaJob(final SourceConnection
         ConfigType.DISCOVER_SCHEMA,
         jobReportingContext,
         source.getSourceDefinitionId(),
-        () -> temporalClient.submitDiscoverSchema(jobId, 0, taskQueue, jobDiscoverCatalogConfig),
+        () -> temporalClient.submitDiscoverSchema(jobId, 0, source.getWorkspaceId(), taskQueue, jobDiscoverCatalogConfig),
         ConnectorJobOutput::getDiscoverCatalogId,
         source.getWorkspaceId(),
         source.getSourceId());

airbyte-commons-server/src/test/java/io/airbyte/commons/server/handlers/JobInputHandlerTest.java

@@ -164,10 +164,11 @@ void testGetSyncWorkflowInput() throws JsonValidationException, ConfigNotFoundEx
     when(stateHandler.getState(new ConnectionIdRequestBody().connectionId(CONNECTION_ID)))
         .thenReturn(new ConnectionState()
             .stateType(ConnectionStateType.LEGACY)
-            .state(STATE.getState()));
+            .state(STATE.getState())
+            .connectionId(CONNECTION_ID));
 
     final JobSyncConfig jobSyncConfig = new JobSyncConfig()
-        .withWorkspaceId(UUID.randomUUID())
+        .withWorkspaceId(WORKSPACE_ID)
         .withDestinationDockerImage("destinationDockerImage")
         .withSourceDockerImage("sourceDockerImage")
         .withConfiguredAirbyteCatalog(mock(ConfiguredAirbyteCatalog.class));
@@ -180,14 +181,14 @@ void testGetSyncWorkflowInput() throws JsonValidationException, ConfigNotFoundEx
     when(job.getScope()).thenReturn(CONNECTION_ID.toString());
 
     final StandardSyncInput expectedStandardSyncInput = new StandardSyncInput()
+        .withConnectionId(CONNECTION_ID)
         .withWorkspaceId(jobSyncConfig.getWorkspaceId())
         .withSourceId(SOURCE_ID)
         .withDestinationId(DESTINATION_ID)
         .withSourceConfiguration(SOURCE_CONFIG_WITH_OAUTH_AND_INJECTED_CONFIG)
         .withDestinationConfiguration(DESTINATION_CONFIG_WITH_OAUTH)
         .withState(STATE)
         .withCatalog(jobSyncConfig.getConfiguredAirbyteCatalog())
-        .withWorkspaceId(jobSyncConfig.getWorkspaceId())
         .withIsReset(false);
 
     final JobRunConfig expectedJobRunConfig = new JobRunConfig()
@@ -197,11 +198,15 @@ void testGetSyncWorkflowInput() throws JsonValidationException, ConfigNotFoundEx
     final IntegrationLauncherConfig expectedSourceLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(JOB_ID))
         .withAttemptId((long) ATTEMPT_NUMBER)
+        .withConnectionId(CONNECTION_ID)
+        .withWorkspaceId(WORKSPACE_ID)
         .withDockerImage(jobSyncConfig.getSourceDockerImage());
 
     final IntegrationLauncherConfig expectedDestinationLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(JOB_ID))
         .withAttemptId((long) ATTEMPT_NUMBER)
+        .withConnectionId(CONNECTION_ID)
+        .withWorkspaceId(jobSyncConfig.getWorkspaceId())
         .withDockerImage(jobSyncConfig.getDestinationDockerImage())
         .withAdditionalEnvironmentVariables(Collections.emptyMap());
 
@@ -235,10 +240,11 @@ void testGetResetSyncWorkflowInput() throws IOException, ApiException, JsonValid
     when(stateHandler.getState(new ConnectionIdRequestBody().connectionId(CONNECTION_ID)))
         .thenReturn(new ConnectionState()
             .stateType(ConnectionStateType.LEGACY)
-            .state(STATE.getState()));
+            .state(STATE.getState())
+            .connectionId(CONNECTION_ID));
 
     final JobResetConnectionConfig jobResetConfig = new JobResetConnectionConfig()
-        .withWorkspaceId(UUID.randomUUID())
+        .withWorkspaceId(WORKSPACE_ID)
         .withDestinationDockerImage("destinationDockerImage")
         .withConfiguredAirbyteCatalog(mock(ConfiguredAirbyteCatalog.class));
 
@@ -250,14 +256,14 @@ void testGetResetSyncWorkflowInput() throws IOException, ApiException, JsonValid
     when(job.getScope()).thenReturn(CONNECTION_ID.toString());
 
     final StandardSyncInput expectedStandardSyncInput = new StandardSyncInput()
+        .withConnectionId(CONNECTION_ID)
         .withWorkspaceId(jobResetConfig.getWorkspaceId())
         .withSourceId(SOURCE_ID)
         .withDestinationId(DESTINATION_ID)
         .withSourceConfiguration(Jsons.emptyObject())
         .withDestinationConfiguration(DESTINATION_CONFIG_WITH_OAUTH)
         .withState(STATE)
         .withCatalog(jobResetConfig.getConfiguredAirbyteCatalog())
-        .withWorkspaceId(jobResetConfig.getWorkspaceId())
         .withWebhookOperationConfigs(jobResetConfig.getWebhookOperationConfigs())
         .withIsReset(true);
 
@@ -268,11 +274,15 @@ void testGetResetSyncWorkflowInput() throws IOException, ApiException, JsonValid
     final IntegrationLauncherConfig expectedSourceLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(JOB_ID))
         .withAttemptId((long) ATTEMPT_NUMBER)
+        .withConnectionId(CONNECTION_ID)
+        .withWorkspaceId(jobResetConfig.getWorkspaceId())
         .withDockerImage(WorkerConstants.RESET_JOB_SOURCE_DOCKER_IMAGE_STUB);
 
     final IntegrationLauncherConfig expectedDestinationLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(JOB_ID))
         .withAttemptId((long) ATTEMPT_NUMBER)
+        .withConnectionId(CONNECTION_ID)
+        .withWorkspaceId(jobResetConfig.getWorkspaceId())
         .withDockerImage(jobResetConfig.getDestinationDockerImage())
         .withAdditionalEnvironmentVariables(Collections.emptyMap());
 
@@ -314,7 +324,7 @@ void testGetCheckConnectionInputs() throws JsonValidationException, ConfigNotFou
         .thenReturn(SOURCE_CONFIG_WITH_OAUTH);
 
     final JobSyncConfig jobSyncConfig = new JobSyncConfig()
-        .withWorkspaceId(UUID.randomUUID())
+        .withWorkspaceId(WORKSPACE_ID)
         .withDestinationDockerImage("destinationDockerImage")
         .withSourceDockerImage("sourceDockerImage")
         .withConfiguredAirbyteCatalog(mock(ConfiguredAirbyteCatalog.class));
@@ -329,11 +339,15 @@ void testGetCheckConnectionInputs() throws JsonValidationException, ConfigNotFou
     final IntegrationLauncherConfig expectedSourceLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(JOB_ID))
         .withAttemptId((long) ATTEMPT_NUMBER)
+        .withWorkspaceId(WORKSPACE_ID)
+        .withConnectionId(CONNECTION_ID)
         .withDockerImage(jobSyncConfig.getSourceDockerImage());
 
     final IntegrationLauncherConfig expectedDestinationLauncherConfig = new IntegrationLauncherConfig()
         .withJobId(String.valueOf(JOB_ID))
         .withAttemptId((long) ATTEMPT_NUMBER)
+        .withWorkspaceId(WORKSPACE_ID)
+        .withConnectionId(CONNECTION_ID)
         .withDockerImage(jobSyncConfig.getDestinationDockerImage())
         .withAdditionalEnvironmentVariables(Collections.emptyMap());
 

airbyte-commons-server/src/test/java/io/airbyte/commons/server/scheduler/DefaultSynchronousSchedulerClientTest.java

@@ -79,10 +79,12 @@ class DefaultSynchronousSchedulerClientTest {
       .put("password", "abc")
       .build());
   private static final SourceConnection SOURCE_CONNECTION = new SourceConnection()
+      .withWorkspaceId(WORKSPACE_ID)
       .withSourceId(UUID1)
       .withSourceDefinitionId(UUID2)
       .withConfiguration(CONFIGURATION);
   private static final DestinationConnection DESTINATION_CONNECTION = new DestinationConnection()
+      .withWorkspaceId(WORKSPACE_ID)
       .withDestinationId(UUID1)
       .withDestinationDefinitionId(UUID2)
       .withConfiguration(CONFIGURATION);
@@ -223,7 +225,7 @@ void testCreateSourceCheckConnectionJob() throws IOException {
 
       final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class);
       final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput);
-      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(CHECK_TASK_QUEUE), eq(jobCheckConnectionConfig)))
+      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(WORKSPACE_ID), eq(CHECK_TASK_QUEUE), eq(jobCheckConnectionConfig)))
           .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
       final SynchronousResponse<StandardCheckConnectionOutput> response =
           schedulerClient.createSourceCheckConnectionJob(SOURCE_CONNECTION, ACTOR_DEFINITION_VERSION, false);
@@ -246,7 +248,7 @@ void testCreateSourceCheckConnectionJobWithConfigInjection() throws IOException
 
       final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class);
       final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput);
-      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(CHECK_TASK_QUEUE), eq(jobCheckConnectionConfig)))
+      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(WORKSPACE_ID), eq(CHECK_TASK_QUEUE), eq(jobCheckConnectionConfig)))
           .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
       final SynchronousResponse<StandardCheckConnectionOutput> response =
           schedulerClient.createSourceCheckConnectionJob(SOURCE_CONNECTION, ACTOR_DEFINITION_VERSION, false);
@@ -265,7 +267,7 @@ void testCreateDestinationCheckConnectionJob() throws IOException {
 
       final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class);
       final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput);
-      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(CHECK_TASK_QUEUE), eq(jobCheckConnectionConfig)))
+      when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(WORKSPACE_ID), eq(CHECK_TASK_QUEUE), eq(jobCheckConnectionConfig)))
           .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
       final SynchronousResponse<StandardCheckConnectionOutput> response =
           schedulerClient.createDestinationCheckConnectionJob(DESTINATION_CONNECTION, ACTOR_DEFINITION_VERSION, false);
@@ -277,8 +279,9 @@ void testCreateDestinationCheckConnectionJob() throws IOException {
     void testCreateDiscoverSchemaJob() throws IOException {
       final UUID expectedCatalogId = UUID.randomUUID();
       final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withDiscoverCatalogId(expectedCatalogId);
-      when(temporalClient.submitDiscoverSchema(any(UUID.class), eq(0), eq(DISCOVER_TASK_QUEUE), any(JobDiscoverCatalogConfig.class)))
-          .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
+      when(
+          temporalClient.submitDiscoverSchema(any(UUID.class), eq(0), eq(WORKSPACE_ID), eq(DISCOVER_TASK_QUEUE), any(JobDiscoverCatalogConfig.class)))
+              .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true)));
       final SynchronousResponse<UUID> response =
           schedulerClient.createDiscoverSchemaJob(SOURCE_CONNECTION, ACTOR_DEFINITION_VERSION, false);
       assertEquals(expectedCatalogId, response.getOutput());

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

@@ -387,12 +387,14 @@ public TemporalResponse<ConnectorJobOutput> submitGetSpec(final UUID jobId, fina
    */
   public TemporalResponse<ConnectorJobOutput> submitCheckConnection(final UUID jobId,
                                                                     final int attempt,
+                                                                    final UUID workspaceId,
                                                                     final String taskQueue,
                                                                     final JobCheckConnectionConfig config) {
     final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt);
     final IntegrationLauncherConfig launcherConfig = new IntegrationLauncherConfig()
         .withJobId(jobId.toString())
         .withAttemptId((long) attempt)
+        .withWorkspaceId(workspaceId)
         .withDockerImage(config.getDockerImage())
         .withProtocolVersion(config.getProtocolVersion())
         .withIsCustomConnector(config.getIsCustomConnector());
@@ -416,12 +418,14 @@ public TemporalResponse<ConnectorJobOutput> submitCheckConnection(final UUID job
    */
   public TemporalResponse<ConnectorJobOutput> submitDiscoverSchema(final UUID jobId,
                                                                    final int attempt,
+                                                                   final UUID workspaceId,
                                                                    final String taskQueue,
                                                                    final JobDiscoverCatalogConfig config) {
     final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt);
     final IntegrationLauncherConfig launcherConfig = new IntegrationLauncherConfig()
         .withJobId(jobId.toString())
         .withAttemptId((long) attempt)
+        .withWorkspaceId(workspaceId)
         .withDockerImage(config.getDockerImage())
         .withProtocolVersion(config.getProtocolVersion())
         .withIsCustomConnector(config.getIsCustomConnector());

airbyte-commons-temporal/src/test/java/io/airbyte/commons/temporal/TemporalClientTest.java

@@ -249,7 +249,7 @@ void testSubmitCheckConnection() {
       final StandardCheckConnectionInput input = new StandardCheckConnectionInput()
           .withConnectionConfiguration(checkConnectionConfig.getConnectionConfiguration());
 
-      temporalClient.submitCheckConnection(JOB_UUID, ATTEMPT_ID, CHECK_TASK_QUEUE, checkConnectionConfig);
+      temporalClient.submitCheckConnection(JOB_UUID, ATTEMPT_ID, WORKSPACE_ID, CHECK_TASK_QUEUE, checkConnectionConfig);
       checkConnectionWorkflow.run(JOB_RUN_CONFIG, UUID_LAUNCHER_CONFIG, input);
       verify(workflowClient).newWorkflowStub(CheckConnectionWorkflow.class,
           TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.CHECK_CONNECTION));
@@ -266,7 +266,7 @@ void testSubmitDiscoverSchema() {
       final StandardDiscoverCatalogInput input = new StandardDiscoverCatalogInput()
           .withConnectionConfiguration(checkConnectionConfig.getConnectionConfiguration());
 
-      temporalClient.submitDiscoverSchema(JOB_UUID, ATTEMPT_ID, DISCOVER_TASK_QUEUE, checkConnectionConfig);
+      temporalClient.submitDiscoverSchema(JOB_UUID, ATTEMPT_ID, WORKSPACE_ID, DISCOVER_TASK_QUEUE, checkConnectionConfig);
       discoverCatalogWorkflow.run(JOB_RUN_CONFIG, UUID_LAUNCHER_CONFIG, input);
       verify(workflowClient).newWorkflowStub(DiscoverCatalogWorkflow.class,
           TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.DISCOVER_SCHEMA));

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

@@ -117,6 +117,8 @@ public boolean transform(final String jobId,
               CUSTOM_STEP,
               jobId,
               attempt,
+              null, // TODO: Provide connectionId
+              null, // TODO: Provide workspaceId
               jobRoot,
               dbtConfig.getDockerImage(),
               false,