pipe through to worker

cgardens · cgardens · commit 511587cc8808 · 2022-02-19T17:00:36.000-08:00
diff --git a/airbyte-config/models/src/main/resources/types/ActorType.yaml b/airbyte-config/models/src/main/resources/types/ActorType.yaml
@@ -0,0 +1,9 @@
+---
+"$schema": http://json-schema.org/draft-07/schema#
+"$id": https://github.com/airbytehq/airbyte/blob/master/airbyte-config/models/src/main/resources/types/ActorType.yaml
+title: ActorType
+description: enum that describes different types of actors
+type: string
+enum:
+  - source
+  - destination
diff --git a/airbyte-config/models/src/main/resources/types/JobSyncConfig.yaml b/airbyte-config/models/src/main/resources/types/JobSyncConfig.yaml
@@ -39,6 +39,10 @@ properties:
   destinationDockerImage:
     description: Image name of the destination with tag.
     type: string
+  sourceResourceRequirements:
+    "$ref": ActorDefinitionResourceRequirements.yaml
+  destinationResourceRequirements:
+    "$ref": ActorDefinitionResourceRequirements.yaml
   operationSequence:
     description: Sequence of configurations of operations to apply as part of the sync
     type: array
diff --git a/airbyte-config/models/src/main/resources/types/StandardSyncInput.yaml b/airbyte-config/models/src/main/resources/types/StandardSyncInput.yaml
@@ -35,11 +35,16 @@ properties:
   catalog:
     description: the configured airbyte catalog
     type: object
+    # necessary because the configuration declaration is in a separate package.
     existingJavaType: io.airbyte.protocol.models.ConfiguredAirbyteCatalog
   state:
     description: optional state of the previous run. this object is defined per integration.
     "$ref": State.yaml
   resourceRequirements:
-    description: optional resource requirements to run sync workers
+    description: optional resource requirements to run sync workers (confusing - overrides source and destination). takes precedence over source and destination resource requirements.
     type: object
-    existingJavaType: io.airbyte.config.ResourceRequirements
+    "$ref": ResourceRequirements.yaml
+  sourceResourceRequirements:
+    "$ref": ActorDefinitionResourceRequirements.yaml
+  destinationResourceRequirements:
+    "$ref": ActorDefinitionResourceRequirements.yaml
diff --git a/airbyte-scheduler/client/src/main/java/io/airbyte/scheduler/client/DefaultSchedulerJobClient.java b/airbyte-scheduler/client/src/main/java/io/airbyte/scheduler/client/DefaultSchedulerJobClient.java
@@ -4,6 +4,7 @@
 
 package io.airbyte.scheduler.client;
 
+import io.airbyte.config.ActorDefinitionResourceRequirements;
 import io.airbyte.config.DestinationConnection;
 import io.airbyte.config.SourceConnection;
 import io.airbyte.config.StandardSync;
@@ -35,15 +36,19 @@ public Job createOrGetActiveSyncJob(final SourceConnection source,
                                       final StandardSync standardSync,
                                       final String sourceDockerImage,
                                       final String destinationDockerImage,
-                                      final List<StandardSyncOperation> standardSyncOperations)
+                                      final List<StandardSyncOperation> standardSyncOperations,
+                                      final ActorDefinitionResourceRequirements sourceResourceRequirements,
+                                      final ActorDefinitionResourceRequirements destinationResourceRequirements)
       throws IOException {
     final Optional<Long> jobIdOptional = jobCreator.createSyncJob(
         source,
         destination,
         standardSync,
         sourceDockerImage,
         destinationDockerImage,
-        standardSyncOperations);
+        standardSyncOperations,
+        sourceResourceRequirements,
+        destinationResourceRequirements);
 
     final long jobId = jobIdOptional.isEmpty()
         ? jobPersistence.getLastReplicationJob(standardSync.getConnectionId()).orElseThrow(() -> new RuntimeException("No job available")).getId()
diff --git a/airbyte-scheduler/client/src/main/java/io/airbyte/scheduler/client/SchedulerJobClient.java b/airbyte-scheduler/client/src/main/java/io/airbyte/scheduler/client/SchedulerJobClient.java
@@ -4,6 +4,7 @@
 
 package io.airbyte.scheduler.client;
 
+import io.airbyte.config.ActorDefinitionResourceRequirements;
 import io.airbyte.config.DestinationConnection;
 import io.airbyte.config.SourceConnection;
 import io.airbyte.config.StandardSync;
@@ -24,7 +25,9 @@ Job createOrGetActiveSyncJob(SourceConnection source,
                                StandardSync standardSync,
                                String sourceDockerImage,
                                String destinationDockerImage,
-                               List<StandardSyncOperation> standardSyncOperations)
+                               List<StandardSyncOperation> standardSyncOperations,
+                               ActorDefinitionResourceRequirements sourceResourceRequirements,
+                               ActorDefinitionResourceRequirements destinationResourceRequirements)
       throws IOException;
 
   Job createOrGetActiveResetConnectionJob(DestinationConnection destination,
diff --git a/airbyte-scheduler/client/src/test/java/io/airbyte/scheduler/client/DefaultSchedulerJobClientTest.java b/airbyte-scheduler/client/src/test/java/io/airbyte/scheduler/client/DefaultSchedulerJobClientTest.java
@@ -46,11 +46,13 @@ void testCreateSyncJob() throws IOException {
     final DestinationConnection destination = mock(DestinationConnection.class);
     final StandardSync standardSync = mock(StandardSync.class);
     final String destinationDockerImage = "airbyte/spaceport";
-    when(jobCreator.createSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of()))
+    when(jobCreator.createSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of(), null, null))
         .thenReturn(Optional.of(JOB_ID));
     when(jobPersistence.getJob(JOB_ID)).thenReturn(job);
 
-    assertEquals(job, client.createOrGetActiveSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of()));
+    assertEquals(
+        job,
+        client.createOrGetActiveSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of(), null, null));
   }
 
   @Test
@@ -61,14 +63,23 @@ void testCreateSyncJobAlreadyExist() throws IOException {
     final UUID connectionUuid = UUID.randomUUID();
     when(standardSync.getConnectionId()).thenReturn(connectionUuid);
     final String destinationDockerImage = "airbyte/spaceport";
-    when(jobCreator.createSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of())).thenReturn(Optional.empty());
+    when(jobCreator.createSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of(), null, null))
+        .thenReturn(Optional.empty());
 
     final Job currentJob = mock(Job.class);
     when(currentJob.getId()).thenReturn(42L);
     when(jobPersistence.getLastReplicationJob(connectionUuid)).thenReturn(Optional.of(currentJob));
     when(jobPersistence.getJob(42L)).thenReturn(currentJob);
 
-    assertEquals(currentJob, client.createOrGetActiveSyncJob(source, destination, standardSync, DOCKER_IMAGE, destinationDockerImage, List.of()));
+    assertEquals(currentJob, client.createOrGetActiveSyncJob(
+        source,
+        destination,
+        standardSync,
+        DOCKER_IMAGE,
+        destinationDockerImage,
+        List.of(),
+        null,
+        null));
   }
 
   @Test
diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java
@@ -4,6 +4,7 @@
 
 package io.airbyte.scheduler.persistence;
 
+import io.airbyte.config.ActorDefinitionResourceRequirements;
 import io.airbyte.config.DestinationConnection;
 import io.airbyte.config.JobConfig;
 import io.airbyte.config.JobConfig.ConfigType;
@@ -41,7 +42,9 @@ public Optional<Long> createSyncJob(final SourceConnection source,
                                       final StandardSync standardSync,
                                       final String sourceDockerImageName,
                                       final String destinationDockerImageName,
-                                      final List<StandardSyncOperation> standardSyncOperations)
+                                      final List<StandardSyncOperation> standardSyncOperations,
+                                      final ActorDefinitionResourceRequirements sourceResourceReqs,
+                                      final ActorDefinitionResourceRequirements destinationResourceReqs)
       throws IOException {
     // reusing this isn't going to quite work.
     final JobSyncConfig jobSyncConfig = new JobSyncConfig()
@@ -55,7 +58,9 @@ public Optional<Long> createSyncJob(final SourceConnection source,
         .withOperationSequence(standardSyncOperations)
         .withConfiguredAirbyteCatalog(standardSync.getCatalog())
         .withState(null)
-        .withResourceRequirements(getJobResourceRequirements(standardSync));
+        .withResourceRequirements(getJobResourceRequirements(standardSync))
+        .withSourceResourceRequirements(sourceResourceReqs)
+        .withDestinationResourceRequirements(destinationResourceReqs);
 
     configRepository.getConnectionState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState);
 
diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobCreator.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobCreator.java
@@ -4,6 +4,7 @@
 
 package io.airbyte.scheduler.persistence;
 
+import io.airbyte.config.ActorDefinitionResourceRequirements;
 import io.airbyte.config.DestinationConnection;
 import io.airbyte.config.SourceConnection;
 import io.airbyte.config.StandardSync;
@@ -29,7 +30,9 @@ Optional<Long> createSyncJob(SourceConnection source,
                                StandardSync standardSync,
                                String sourceDockerImage,
                                String destinationDockerImage,
-                               List<StandardSyncOperation> standardSyncOperations)
+                               List<StandardSyncOperation> standardSyncOperations,
+                               ActorDefinitionResourceRequirements sourceResourceReqs,
+                               ActorDefinitionResourceRequirements destinationResourceReqs)
       throws IOException;
 
   /**
diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/job_factory/DefaultSyncJobFactory.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/job_factory/DefaultSyncJobFactory.java
@@ -50,9 +50,10 @@ public Long create(final UUID connectionId) {
           destinationConnection.getWorkspaceId(),
           destinationConnection.getConfiguration());
       destinationConnection.withConfiguration(destinationConfiguration);
-      final StandardSourceDefinition sourceDefinition = configRepository.getStandardSourceDefinition(sourceConnection.getSourceDefinitionId());
-      final StandardDestinationDefinition destinationDefinition =
-          configRepository.getStandardDestinationDefinition(destinationConnection.getDestinationDefinitionId());
+      final StandardSourceDefinition sourceDefinition = configRepository
+          .getStandardSourceDefinition(sourceConnection.getSourceDefinitionId());
+      final StandardDestinationDefinition destinationDefinition = configRepository
+          .getStandardDestinationDefinition(destinationConnection.getDestinationDefinitionId());
 
       final String sourceImageName = DockerUtils.getTaggedImageName(sourceDefinition.getDockerRepository(), sourceDefinition.getDockerImageTag());
       final String destinationImageName =
@@ -70,7 +71,9 @@ public Long create(final UUID connectionId) {
           standardSync,
           sourceImageName,
           destinationImageName,
-          standardSyncOperations)
+          standardSyncOperations,
+          sourceDefinition.getResourceRequirements(),
+          destinationDefinition.getResourceRequirements())
           .orElseThrow(() -> new IllegalStateException("We shouldn't be trying to create a new sync job if there is one running already."));
 
     } catch (final IOException | JsonValidationException | ConfigNotFoundException e) {
diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobCreatorTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobCreatorTest.java
@@ -151,7 +151,9 @@ void testCreateSyncJob() throws IOException {
         STANDARD_SYNC,
         SOURCE_IMAGE_NAME,
         DESTINATION_IMAGE_NAME,
-        List.of(STANDARD_SYNC_OPERATION)).orElseThrow();
+        List.of(STANDARD_SYNC_OPERATION),
+        null,
+        null).orElseThrow();
     assertEquals(JOB_ID, jobId);
   }
 
@@ -182,7 +184,9 @@ void testCreateSyncJobEnsureNoQueuing() throws IOException {
         STANDARD_SYNC,
         SOURCE_IMAGE_NAME,
         DESTINATION_IMAGE_NAME,
-        List.of(STANDARD_SYNC_OPERATION)).isEmpty());
+        List.of(STANDARD_SYNC_OPERATION),
+        null,
+        null).isEmpty());
   }
 
   @Test
diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/job_factory/DefaultSyncJobFactoryTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/job_factory/DefaultSyncJobFactoryTest.java
@@ -62,7 +62,7 @@ void createSyncJobFromConnectionId() throws JsonValidationException, ConfigNotFo
     when(configRepository.getSourceConnection(sourceId)).thenReturn(sourceConnection);
     when(configRepository.getDestinationConnection(destinationId)).thenReturn(destinationConnection);
     when(configRepository.getStandardSyncOperation(operationId)).thenReturn(operation);
-    when(jobCreator.createSyncJob(sourceConnection, destinationConnection, standardSync, srcDockerImage, dstDockerImage, operations))
+    when(jobCreator.createSyncJob(sourceConnection, destinationConnection, standardSync, srcDockerImage, dstDockerImage, operations, null, null))
         .thenReturn(Optional.of(jobId));
     when(configRepository.getStandardSourceDefinition(sourceDefinitionId))
         .thenReturn(new StandardSourceDefinition().withSourceDefinitionId(sourceDefinitionId).withDockerRepository(srcDockerRepo)
@@ -77,7 +77,7 @@ void createSyncJobFromConnectionId() throws JsonValidationException, ConfigNotFo
     assertEquals(jobId, actualJobId);
 
     verify(jobCreator)
-        .createSyncJob(sourceConnection, destinationConnection, standardSync, srcDockerImage, dstDockerImage, operations);
+        .createSyncJob(sourceConnection, destinationConnection, standardSync, srcDockerImage, dstDockerImage, operations, null, null);
   }
 
 }
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/helper/ResourceRequirementsUtils.java b/airbyte-workers/src/main/java/io/airbyte/workers/helper/ResourceRequirementsUtils.java
@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.workers.helper;
+
+import com.google.common.base.Preconditions;
+import io.airbyte.config.ActorDefinitionResourceRequirements;
+import io.airbyte.config.ActorType;
+import io.airbyte.config.JobTypeResourceLimit;
+import io.airbyte.config.JobTypeResourceLimit.JobType;
+import io.airbyte.config.ResourceRequirements;
+import io.airbyte.config.StandardSyncInput;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+public class ResourceRequirementsUtils {
+
+  /**
+   * Given an actor type and job type, returns resource requirements, if present. Otherwise, null.
+   * Sync-level resource requirements, take precedence over connector definition level ones.
+   *
+   * @param input - sync input
+   * @param actorType - type of actor to extract resource requirements for
+   * @param jobType - toe of job to extract resource requirements for
+   * @return resource requirements, if present, otherwise null.
+   */
+  public static ResourceRequirements getResourceRequirements(final StandardSyncInput input, final ActorType actorType, final JobType jobType) {
+    final ActorDefinitionResourceRequirements actorDefResourceReqs = getResourceRequirementsForActorType(input, actorType);
+    if (input.getResourceRequirements() != null) {
+      return input.getResourceRequirements();
+    } else if (actorDefResourceReqs != null) {
+      final Optional<ResourceRequirements> resourceRequirementsForJobType = getResourceRequirementsForJobType(actorDefResourceReqs, jobType);
+      return resourceRequirementsForJobType.orElse(actorDefResourceReqs.getDefault()); // default can return null.
+    } else {
+      return null;
+    }
+
+  }
+
+  private static Optional<ResourceRequirements> getResourceRequirementsForJobType(final ActorDefinitionResourceRequirements actorDefResourceReqs,
+                                                                                  final JobType jobType) {
+    final List<ResourceRequirements> jobTypeResourceRequirement = actorDefResourceReqs.getJobSpecific()
+        .stream()
+        .filter(jobSpecific -> jobSpecific.getJobType() == jobType).map(JobTypeResourceLimit::getResourceRequirements).collect(
+            Collectors.toList());
+
+    Preconditions.checkArgument(jobTypeResourceRequirement.size() <= 1, "Should only have one resource requirement per job type.");
+    return jobTypeResourceRequirement.isEmpty() ? Optional.empty() : Optional.of(jobTypeResourceRequirement.get(0));
+  }
+
+  private static ActorDefinitionResourceRequirements getResourceRequirementsForActorType(final StandardSyncInput input, final ActorType actorType) {
+    switch (actorType) {
+      case SOURCE -> {
+        return input.getSourceResourceRequirements();
+      }
+      case DESTINATION -> {
+        return input.getDestinationResourceRequirements();
+      }
+      default -> {
+        throw new IllegalArgumentException("Unrecognized actor type.");
+      }
+    }
+  }
+
+}
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/GenerateInputActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/GenerateInputActivityImpl.java
@@ -65,7 +65,9 @@ public GeneratedJobInput getSyncWorkflowInput(final SyncInput input) {
           .withOperationSequence(config.getOperationSequence())
           .withCatalog(config.getConfiguredAirbyteCatalog())
           .withState(config.getState())
-          .withResourceRequirements(config.getResourceRequirements());
+          .withResourceRequirements(config.getResourceRequirements())
+          .withSourceResourceRequirements(config.getSourceResourceRequirements())
+          .withDestinationResourceRequirements(config.getDestinationResourceRequirements());
 
       return new GeneratedJobInput(jobRunConfig, sourceLauncherConfig, destinationLauncherConfig, syncInput);
 
diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationActivityImpl.java
@@ -7,9 +7,11 @@
 import com.google.common.annotations.VisibleForTesting;
 import io.airbyte.commons.functional.CheckedSupplier;
 import io.airbyte.commons.json.Jsons;
+import io.airbyte.config.ActorType;
 import io.airbyte.config.AirbyteConfigValidator;
 import io.airbyte.config.ConfigSchema;
 import io.airbyte.config.Configs.WorkerEnvironment;
+import io.airbyte.config.JobTypeResourceLimit.JobType;
 import io.airbyte.config.ReplicationOutput;
 import io.airbyte.config.StandardSyncInput;
 import io.airbyte.config.StandardSyncOutput;
@@ -24,6 +26,7 @@
 import io.airbyte.workers.WorkerApp;
 import io.airbyte.workers.WorkerConfigs;
 import io.airbyte.workers.WorkerConstants;
+import io.airbyte.workers.helper.ResourceRequirementsUtils;
 import io.airbyte.workers.process.AirbyteIntegrationLauncher;
 import io.airbyte.workers.process.IntegrationLauncher;
 import io.airbyte.workers.process.ProcessFactory;
@@ -180,13 +183,13 @@ private CheckedSupplier<Worker<StandardSyncInput, ReplicationOutput>, Exception>
           Math.toIntExact(sourceLauncherConfig.getAttemptId()),
           sourceLauncherConfig.getDockerImage(),
           processFactory,
-          syncInput.getResourceRequirements());
+          ResourceRequirementsUtils.getResourceRequirements(syncInput, ActorType.SOURCE, JobType.SYNC));
       final IntegrationLauncher destinationLauncher = new AirbyteIntegrationLauncher(
           destinationLauncherConfig.getJobId(),
           Math.toIntExact(destinationLauncherConfig.getAttemptId()),
           destinationLauncherConfig.getDockerImage(),
           processFactory,
-          syncInput.getResourceRequirements());
+          ResourceRequirementsUtils.getResourceRequirements(syncInput, ActorType.DESTINATION, JobType.SYNC));
 
       // reset jobs use an empty source to induce resetting all data in destination.
       final AirbyteSource airbyteSource =
