From 950269b49220c2be47a47eb068e7bbe4ef08e7e8 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Mon, 21 Nov 2022 11:27:57 -0800 Subject: [PATCH 01/18] kubeprocessfactory flag --- .../src/main/java/io/airbyte/workers/WorkerConfigs.java | 8 ++++++++ .../workers/process/AirbyteIntegrationLauncher.java | 9 +++++++++ .../io/airbyte/workers/process/DockerProcessFactory.java | 1 + .../io/airbyte/workers/process/KubeProcessFactory.java | 7 ++++++- .../java/io/airbyte/workers/process/ProcessFactory.java | 2 ++ .../src/main/java/io/airbyte/config/Configs.java | 5 +++++ .../src/main/java/io/airbyte/config/EnvConfigs.java | 6 ++++++ 7 files changed, 37 insertions(+), 1 deletion(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java index a8997e271088..71dee3125214 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java @@ -18,6 +18,7 @@ public class WorkerConfigs { private final ResourceRequirements resourceRequirements; private final List workerKubeTolerations; private final Map workerKubeNodeSelectors; + private final Map workerIsolatedKubeNodeSelectors; private final Map workerKubeAnnotations; private final String jobImagePullSecret; private final String jobImagePullPolicy; @@ -41,6 +42,7 @@ public WorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), configs.getJobKubeNodeSelectors(), + configs.getIsolatedJobKubeNodeSelectors(), configs.getJobKubeAnnotations(), configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -72,6 +74,7 @@ public static WorkerConfigs buildSpecWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, + configs.getIsolatedJobKubeNodeSelectors(), annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -103,6 +106,7 @@ public static WorkerConfigs buildCheckWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getCheckJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, + configs.getIsolatedJobKubeNodeSelectors(), annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -134,6 +138,7 @@ public static WorkerConfigs buildDiscoverWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, + configs.getIsolatedJobKubeNodeSelectors(), annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -154,6 +159,7 @@ public static WorkerConfigs buildReplicationWorkerConfigs(final Configs configs) .withMemoryLimit(configs.getReplicationOrchestratorMemoryLimit()), configs.getJobKubeTolerations(), configs.getJobKubeNodeSelectors(), + configs.getIsolatedJobKubeNodeSelectors(), configs.getJobKubeAnnotations(), configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -180,6 +186,8 @@ public Map getworkerKubeNodeSelectors() { return workerKubeNodeSelectors; } + public Map getWorkerIsolatedKubeNodeSelectors() {return workerIsolatedKubeNodeSelectors;} + public Map getWorkerKubeAnnotations() { return workerKubeAnnotations; } diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java index 7f8542f9b116..d2352b054d70 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java @@ -70,6 +70,10 @@ public AirbyteIntegrationLauncher(final String jobId, this.featureFlags = new EnvVariableFeatureFlags(); } + public boolean doesProcessNeedToBeIsolated() { + return false; + } + @Trace(operationName = WORKER_OPERATION_NAME) @Override public Process spec(final Path jobRoot) throws WorkerException { @@ -81,6 +85,7 @@ public Process spec(final Path jobRoot) throws WorkerException { jobRoot, imageName, false, + false, Collections.emptyMap(), null, resourceRequirement, @@ -101,6 +106,7 @@ public Process check(final Path jobRoot, final String configFilename, final Stri jobRoot, imageName, false, + false, ImmutableMap.of(configFilename, configContents), null, resourceRequirement, @@ -122,6 +128,7 @@ public Process discover(final Path jobRoot, final String configFilename, final S jobRoot, imageName, false, + false, ImmutableMap.of(configFilename, configContents), null, resourceRequirement, @@ -167,6 +174,7 @@ public Process read(final Path jobRoot, jobRoot, imageName, false, + false, files, null, resourceRequirement, @@ -195,6 +203,7 @@ public Process write(final Path jobRoot, attempt, jobRoot, imageName, + false, true, files, null, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/DockerProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/DockerProcessFactory.java index 743a429a772d..5ed26798030b 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/DockerProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/DockerProcessFactory.java @@ -83,6 +83,7 @@ public Process create(final String jobType, final int attempt, final Path jobRoot, final String imageName, + final boolean usesIsolatedPool, final boolean usesStdin, final Map files, final String entrypoint, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index 99227c82cc8f..6bce335ac17a 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -37,6 +37,8 @@ public class KubeProcessFactory implements ProcessFactory { private final String processRunnerHost; private final boolean isOrchestrator; + private boolean useIsolatedNodePool; + /** * Sets up a process factory with the default processRunnerHost. */ @@ -85,6 +87,7 @@ public Process create( final int attempt, final Path jobRoot, final String imageName, + final boolean usesIsolatedPool, final boolean usesStdin, final Map files, final String entrypoint, @@ -107,6 +110,8 @@ public Process create( final var allLabels = getLabels(jobId, attempt, customLabels); + final var nodeSelectors = usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors() : workerConfigs.getworkerKubeNodeSelectors(); + return new KubePodProcess( isOrchestrator, processRunnerHost, @@ -125,7 +130,7 @@ public Process create( resourceRequirements, workerConfigs.getJobImagePullSecret(), workerConfigs.getWorkerKubeTolerations(), - workerConfigs.getworkerKubeNodeSelectors(), + nodeSelectors, allLabels, workerConfigs.getWorkerKubeAnnotations(), workerConfigs.getJobSocatImage(), diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/ProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/ProcessFactory.java index 676114796ca6..d9cf45084a06 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/ProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/ProcessFactory.java @@ -26,6 +26,7 @@ public interface ProcessFactory { * @param attempt attempt Id * @param jobPath Workspace directory to run the process from. * @param imageName Docker image name to start the process from. + * @param usesIsolatedPool whether to use isolated pool to run the jobs. * @param files File name to contents map that will be written into the working dir of the process * prior to execution. * @param entrypoint If not null, the default entrypoint program of the docker image can be changed @@ -43,6 +44,7 @@ Process create(String jobType, int attempt, final Path jobPath, final String imageName, + final boolean usesIsolatedPool, final boolean usesStdin, final Map files, final String entrypoint, diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java b/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java index 35337c1b738f..51b8f8c4a8c4 100644 --- a/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java +++ b/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java @@ -405,6 +405,11 @@ public interface Configs { */ Map getJobKubeNodeSelectors(); + /** + * Define an isolated kube node selectors, so we can run risky images in it. + */ + Map getIsolatedJobKubeNodeSelectors(); + /** * Define node selectors for Spec job pods specifically. Each kv-pair is separated by a `,`. */ diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java b/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java index 676c06c9fe58..9928ec62b987 100644 --- a/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java +++ b/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java @@ -72,6 +72,7 @@ public class EnvConfigs implements Configs { public static final String JOB_KUBE_SIDECAR_CONTAINER_IMAGE_PULL_POLICY = "JOB_KUBE_SIDECAR_CONTAINER_IMAGE_PULL_POLICY"; public static final String JOB_KUBE_TOLERATIONS = "JOB_KUBE_TOLERATIONS"; public static final String JOB_KUBE_NODE_SELECTORS = "JOB_KUBE_NODE_SELECTORS"; + public static final String JOB_ISOLATED_KUBE_NODE_SELECTORS = "JOB_ISOLATED_KUBE_NODE_SELECTORS"; public static final String JOB_KUBE_ANNOTATIONS = "JOB_KUBE_ANNOTATIONS"; public static final String JOB_KUBE_SOCAT_IMAGE = "JOB_KUBE_SOCAT_IMAGE"; public static final String JOB_KUBE_BUSYBOX_IMAGE = "JOB_KUBE_BUSYBOX_IMAGE"; @@ -606,6 +607,11 @@ public Map getJobKubeNodeSelectors() { return splitKVPairsFromEnvString(getEnvOrDefault(JOB_KUBE_NODE_SELECTORS, "")); } + @Override + public Map getIsolatedJobKubeNodeSelectors() { + return splitKVPairsFromEnvString(getEnvOrDefault(JOB_ISOLATED_KUBE_NODE_SELECTORS, "")); + } + /** * Returns a map of node selectors for Spec job pods specifically. * From 8c77be55cd34e9a6b6d4f7e4146fa391aba756e0 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 22 Nov 2022 12:04:54 -0800 Subject: [PATCH 02/18] plumbing through custom connector bit --- .../commons/temporal/TemporalClient.java | 3 +- .../general/DbtTransformationRunner.java | 2 +- .../workers/helper/EntrypointEnvChecker.java | 1 + .../DefaultNormalizationRunner.java | 2 + .../process/AirbyteIntegrationLauncher.java | 19 +++---- .../DefaultNormalizationRunnerTest.java | 2 +- .../AirbyteIntegrationLauncherTest.java | 12 ++--- .../process/DockerProcessFactoryTest.java | 4 +- .../types/JobCheckConnectionConfig.yaml | 3 ++ .../types/JobDiscoverCatalogConfig.yaml | 3 ++ .../resources/types/JobGetSpecConfig.yaml | 3 ++ .../types/JobResetConnectionConfig.yaml | 6 +++ .../main/resources/types/JobSyncConfig.yaml | 6 +++ .../ReplicationJobOrchestrator.java | 6 ++- .../DestinationDefinitionsHandler.java | 8 +-- .../server/handlers/SchedulerHandler.java | 19 ++++--- .../handlers/SourceDefinitionsHandler.java | 8 +-- .../DefaultSynchronousSchedulerClient.java | 22 +++++--- .../scheduler/SynchronousSchedulerClient.java | 8 +-- .../DestinationDefinitionsHandlerTest.java | 24 ++++----- .../server/handlers/SchedulerHandlerTest.java | 52 +++++++++---------- .../SourceDefinitionsHandlerTest.java | 24 ++++----- ...DefaultSynchronousSchedulerClientTest.java | 16 +++--- .../IntegrationLauncherConfig.yaml | 2 + .../WorkerConfigurationBeanFactory.java | 22 ++++++-- .../CheckConnectionActivityImpl.java | 3 +- .../catalog/DiscoverCatalogActivityImpl.java | 2 +- .../activities/GenerateInputActivityImpl.java | 10 ++-- .../temporal/spec/SpecActivityImpl.java | 5 +- .../sync/ReplicationActivityImpl.java | 6 ++- .../KubePodProcessIntegrationTest.java | 2 +- 31 files changed, 186 insertions(+), 119 deletions(-) diff --git a/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java b/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java index 9abceab23c75..8073ff001cb8 100644 --- a/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java +++ b/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java @@ -332,7 +332,8 @@ public TemporalResponse submitGetSpec(final UUID jobId, fina final IntegrationLauncherConfig launcherConfig = new IntegrationLauncherConfig() .withJobId(jobId.toString()) .withAttemptId((long) attempt) - .withDockerImage(config.getDockerImage()); + .withDockerImage(config.getDockerImage()) + .withIsCustomConnector(config.getIsCustomConnector()); return execute(jobRunConfig, () -> getWorkflowStub(SpecWorkflow.class, TemporalJobType.GET_SPEC).run(jobRunConfig, launcherConfig)); diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java index df0ff0807151..708ce9f7a0a3 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java @@ -98,7 +98,7 @@ public boolean transform(final String jobId, attempt, jobRoot, dbtConfig.getDockerImage(), - false, + false, false, files, "/bin/bash", resourceRequirements, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/helper/EntrypointEnvChecker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/helper/EntrypointEnvChecker.java index 9b59867ada27..a88ae8a62f51 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/helper/EntrypointEnvChecker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/helper/EntrypointEnvChecker.java @@ -41,6 +41,7 @@ public static String getEntrypointEnvVariable(final ProcessFactory processFactor jobRoot, imageName, false, + false, Collections.emptyMap(), "printenv", null, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/normalization/DefaultNormalizationRunner.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/normalization/DefaultNormalizationRunner.java index 3db73387fec0..e7b84e16d825 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/normalization/DefaultNormalizationRunner.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/normalization/DefaultNormalizationRunner.java @@ -138,6 +138,8 @@ private boolean runProcess(final String jobId, attempt, jobRoot, normalizationImageName, + // custom connector does not use normalization + false, false, files, null, resourceRequirements, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java index d2352b054d70..6cf05cf952fa 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java @@ -57,22 +57,23 @@ public class AirbyteIntegrationLauncher implements IntegrationLauncher { private final ResourceRequirements resourceRequirement; private final FeatureFlags featureFlags; + private boolean useIsolatedNodePool; + public AirbyteIntegrationLauncher(final String jobId, final int attempt, final String imageName, final ProcessFactory processFactory, - final ResourceRequirements resourceRequirement) { + final ResourceRequirements resourceRequirement, + final boolean useIsolatedNodePool) { this.jobId = jobId; this.attempt = attempt; this.imageName = imageName; this.processFactory = processFactory; this.resourceRequirement = resourceRequirement; this.featureFlags = new EnvVariableFeatureFlags(); + this.useIsolatedNodePool = useIsolatedNodePool; } - public boolean doesProcessNeedToBeIsolated() { - return false; - } @Trace(operationName = WORKER_OPERATION_NAME) @Override @@ -84,7 +85,7 @@ public Process spec(final Path jobRoot) throws WorkerException { attempt, jobRoot, imageName, - false, + useIsolatedNodePool, false, Collections.emptyMap(), null, @@ -105,7 +106,7 @@ public Process check(final Path jobRoot, final String configFilename, final Stri attempt, jobRoot, imageName, - false, + useIsolatedNodePool, false, ImmutableMap.of(configFilename, configContents), null, @@ -127,7 +128,7 @@ public Process discover(final Path jobRoot, final String configFilename, final S attempt, jobRoot, imageName, - false, + useIsolatedNodePool, false, ImmutableMap.of(configFilename, configContents), null, @@ -173,7 +174,7 @@ public Process read(final Path jobRoot, attempt, jobRoot, imageName, - false, + useIsolatedNodePool, false, files, null, @@ -203,7 +204,7 @@ public Process write(final Path jobRoot, attempt, jobRoot, imageName, - false, + useIsolatedNodePool, true, files, null, diff --git a/airbyte-commons-worker/src/test/java/io/airbyte/workers/normalization/DefaultNormalizationRunnerTest.java b/airbyte-commons-worker/src/test/java/io/airbyte/workers/normalization/DefaultNormalizationRunnerTest.java index 6b182f2c656e..63f798bf141c 100644 --- a/airbyte-commons-worker/src/test/java/io/airbyte/workers/normalization/DefaultNormalizationRunnerTest.java +++ b/airbyte-commons-worker/src/test/java/io/airbyte/workers/normalization/DefaultNormalizationRunnerTest.java @@ -82,7 +82,7 @@ void setup() throws IOException, WorkerException { WorkerConstants.DESTINATION_CATALOG_JSON_FILENAME, Jsons.serialize(catalog)); when(processFactory.create(AirbyteIntegrationLauncher.NORMALIZE_STEP, JOB_ID, JOB_ATTEMPT, jobRoot, - NormalizationRunnerFactory.BASE_NORMALIZATION_IMAGE_NAME, false, files, null, + NormalizationRunnerFactory.BASE_NORMALIZATION_IMAGE_NAME, false, false, files, null, workerConfigs.getResourceRequirements(), Map.of(AirbyteIntegrationLauncher.JOB_TYPE, AirbyteIntegrationLauncher.SYNC_JOB, AirbyteIntegrationLauncher.SYNC_STEP, AirbyteIntegrationLauncher.NORMALIZE_STEP), diff --git a/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java b/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java index 03bf563792b3..a67170d82982 100644 --- a/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java +++ b/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java @@ -63,14 +63,14 @@ class AirbyteIntegrationLauncherTest { @BeforeEach void setUp() { workerConfigs = new WorkerConfigs(new EnvConfigs()); - launcher = new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, FAKE_IMAGE, processFactory, workerConfigs.getResourceRequirements()); + launcher = new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, FAKE_IMAGE, processFactory, workerConfigs.getResourceRequirements(), false); } @Test void spec() throws WorkerException { launcher.spec(JOB_ROOT); - Mockito.verify(processFactory).create(SPEC_JOB, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, Collections.emptyMap(), null, + Mockito.verify(processFactory).create(SPEC_JOB, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, false, Collections.emptyMap(), null, workerConfigs.getResourceRequirements(), Map.of(AirbyteIntegrationLauncher.JOB_TYPE, AirbyteIntegrationLauncher.SPEC_JOB), JOB_METADATA, Map.of(), "spec"); @@ -80,7 +80,7 @@ void spec() throws WorkerException { void check() throws WorkerException { launcher.check(JOB_ROOT, CONFIG, "{}"); - Mockito.verify(processFactory).create(CHECK_JOB, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, CONFIG_FILES, null, + Mockito.verify(processFactory).create(CHECK_JOB, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, false, CONFIG_FILES, null, workerConfigs.getResourceRequirements(), Map.of(JOB_TYPE, CHECK_JOB), JOB_METADATA, @@ -93,7 +93,7 @@ void check() throws WorkerException { void discover() throws WorkerException { launcher.discover(JOB_ROOT, CONFIG, "{}"); - Mockito.verify(processFactory).create(DISCOVER_JOB, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, CONFIG_FILES, null, + Mockito.verify(processFactory).create(DISCOVER_JOB, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, false, CONFIG_FILES, null, workerConfigs.getResourceRequirements(), Map.of(JOB_TYPE, DISCOVER_JOB), JOB_METADATA, @@ -106,7 +106,7 @@ void discover() throws WorkerException { void read() throws WorkerException { launcher.read(JOB_ROOT, CONFIG, "{}", CATALOG, "{}", "state", "{}"); - Mockito.verify(processFactory).create(READ_STEP, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, CONFIG_CATALOG_STATE_FILES, null, + Mockito.verify(processFactory).create(READ_STEP, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, false, CONFIG_CATALOG_STATE_FILES, null, workerConfigs.getResourceRequirements(), Map.of(JOB_TYPE, SYNC_JOB, SYNC_STEP, READ_STEP), JOB_METADATA, @@ -122,7 +122,7 @@ void read() throws WorkerException { void write() throws WorkerException { launcher.write(JOB_ROOT, CONFIG, "{}", CATALOG, "{}"); - Mockito.verify(processFactory).create(WRITE_STEP, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, true, CONFIG_CATALOG_FILES, null, + Mockito.verify(processFactory).create(WRITE_STEP, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, true, CONFIG_CATALOG_FILES, null, workerConfigs.getResourceRequirements(), Map.of(JOB_TYPE, SYNC_JOB, SYNC_STEP, WRITE_STEP), JOB_METADATA, diff --git a/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/DockerProcessFactoryTest.java b/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/DockerProcessFactoryTest.java index ba8170320962..215fb0bb2015 100644 --- a/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/DockerProcessFactoryTest.java +++ b/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/DockerProcessFactoryTest.java @@ -85,7 +85,7 @@ void testFileWriting() throws IOException, WorkerException { final DockerProcessFactory processFactory = new DockerProcessFactory(new WorkerConfigs(new EnvConfigs()), workspaceRoot, null, null, null); - processFactory.create("tester", "job_id", 0, jobRoot, BUSYBOX, false, ImmutableMap.of("config.json", "{\"data\": 2}"), "echo hi", + processFactory.create("tester", "job_id", 0, jobRoot, BUSYBOX, false, false, ImmutableMap.of("config.json", "{\"data\": 2}"), "echo hi", new WorkerConfigs(new EnvConfigs()).getResourceRequirements(), Map.of(), Map.of(), Map.of()); assertEquals( @@ -121,6 +121,7 @@ void testEnvMapSet() throws IOException, WorkerException, InterruptedException { jobRoot, BUSYBOX, false, + false, Map.of(), "/bin/sh", workerConfigs.getResourceRequirements(), @@ -153,6 +154,7 @@ private void waitForDockerToInitialize(final ProcessFactory processFactory, fina jobRoot, BUSYBOX, false, + false, Map.of(), "/bin/sh", workerConfigs.getResourceRequirements(), diff --git a/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml index 7f239733d589..818cd8d24b18 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml @@ -18,3 +18,6 @@ properties: protocolVersion: type: object existingJavaType: io.airbyte.commons.version.Version + isCustomConnector: + description: determine if the running image is a custom connector. + type: boolean \ No newline at end of file diff --git a/airbyte-config/config-models/src/main/resources/types/JobDiscoverCatalogConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobDiscoverCatalogConfig.yaml index a9f189714b69..8f60458a3e11 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobDiscoverCatalogConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobDiscoverCatalogConfig.yaml @@ -33,3 +33,6 @@ properties: protocolVersion: type: object existingJavaType: io.airbyte.commons.version.Version + isCustomConnector: + description: determine if the running image is a custom connector. + type: boolean diff --git a/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml index 0c5bd6fb7bc2..507dfafd9ee6 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml @@ -10,3 +10,6 @@ required: properties: dockerImage: type: string + isCustomConnector: + description: determine if the running image is a custom connector. + type: boolean \ No newline at end of file diff --git a/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml index 160cdd4067d7..6c2529fad220 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml @@ -52,3 +52,9 @@ properties: state: description: optional current state of the connection "$ref": State.yaml + isSourceCustomConnector: + description: determine if the running image of the source is a custom connector. + type: boolean + isDestinationCustomConnector: + description: determine if the running image of the destination is a custom connector. + type: boolean \ No newline at end of file diff --git a/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml index f6b9af129a08..7bb14d29176d 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml @@ -71,3 +71,9 @@ properties: type: object description: optional resource requirements to run sync workers - this is used for containers other than the source/dest containers existingJavaType: io.airbyte.config.ResourceRequirements + isSourceCustomConnector: + description: determine if the source running image is a custom connector. + type: boolean + isDestinationCustomConnector: + description: determine if the destination running image is a custom connector. + type: boolean \ No newline at end of file diff --git a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java index 11580c575faa..6c49bc3c9183 100644 --- a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java +++ b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java @@ -106,7 +106,8 @@ public Optional runJob() throws Exception { Math.toIntExact(sourceLauncherConfig.getAttemptId()), sourceLauncherConfig.getDockerImage(), processFactory, - syncInput.getSourceResourceRequirements()); + syncInput.getSourceResourceRequirements(), + sourceLauncherConfig.getIsCustomConnector()); log.info("Setting up destination launcher..."); final IntegrationLauncher destinationLauncher = new AirbyteIntegrationLauncher( @@ -114,7 +115,8 @@ public Optional runJob() throws Exception { Math.toIntExact(destinationLauncherConfig.getAttemptId()), destinationLauncherConfig.getDockerImage(), processFactory, - syncInput.getDestinationResourceRequirements()); + syncInput.getDestinationResourceRequirements(), + destinationLauncherConfig.getIsCustomConnector()); log.info("Setting up source..."); // reset jobs use an empty source to induce resetting all data in destination. diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java index 6eda6d4e1da5..4223077519a5 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java @@ -216,7 +216,7 @@ public DestinationDefinitionRead createCustomDestinationDefinition(final CustomD private StandardDestinationDefinition destinationDefinitionFromCreate(final DestinationDefinitionCreate destinationDefCreate) throws IOException { final ConnectorSpecification spec = getSpecForImage( destinationDefCreate.getDockerRepository(), - destinationDefCreate.getDockerImageTag()); + destinationDefCreate.getDockerImageTag(), true); final Version airbyteProtocolVersion = AirbyteProtocolVersion.getWithDefault(spec.getProtocolVersion()); @@ -246,7 +246,7 @@ public DestinationDefinitionRead updateDestinationDefinition(final DestinationDe final boolean specNeedsUpdate = !currentDestination.getDockerImageTag().equals(destinationDefinitionUpdate.getDockerImageTag()) || destinationDefinitionUpdate.getDockerImageTag().equals(DEV_IMAGE_TAG); final ConnectorSpecification spec = specNeedsUpdate - ? getSpecForImage(currentDestination.getDockerRepository(), destinationDefinitionUpdate.getDockerImageTag()) + ? getSpecForImage(currentDestination.getDockerRepository(), destinationDefinitionUpdate.getDockerImageTag(), currentDestination.getCustom()) : currentDestination.getSpec(); final ActorDefinitionResourceRequirements updatedResourceReqs = destinationDefinitionUpdate.getResourceRequirements() != null ? ApiPojoConverters.actorDefResourceReqsToInternal(destinationDefinitionUpdate.getResourceRequirements()) @@ -316,9 +316,9 @@ public void deleteCustomDestinationDefinition(final DestinationDefinitionIdWithW deleteDestinationDefinition(new DestinationDefinitionIdRequestBody().destinationDefinitionId(definitionId)); } - private ConnectorSpecification getSpecForImage(final String dockerRepository, final String imageTag) throws IOException { + private ConnectorSpecification getSpecForImage(final String dockerRepository, final String imageTag, boolean isCustomConnector) throws IOException { final String imageName = DockerUtils.getTaggedImageName(dockerRepository, imageTag); - final SynchronousResponse getSpecResponse = schedulerSynchronousClient.createGetSpecJob(imageName); + final SynchronousResponse getSpecResponse = schedulerSynchronousClient.createGetSpecJob(imageName, isCustomConnector); return SpecFetcher.getSpecFromJob(getSpecResponse); } diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index 49a4a75940e9..db7683c390d7 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -143,9 +143,10 @@ public CheckConnectionRead checkSourceConnectionFromSourceId(final SourceIdReque final SourceConnection source = configRepository.getSourceConnection(sourceIdRequestBody.getSourceId()); final StandardSourceDefinition sourceDef = configRepository.getStandardSourceDefinition(source.getSourceDefinitionId()); final String imageName = DockerUtils.getTaggedImageName(sourceDef.getDockerRepository(), sourceDef.getDockerImageTag()); + final boolean isCustomConnector = sourceDef.getCustom(); final Version protocolVersion = new Version(sourceDef.getProtocolVersion()); - return reportConnectionStatus(synchronousSchedulerClient.createSourceCheckConnectionJob(source, imageName, protocolVersion)); + return reportConnectionStatus(synchronousSchedulerClient.createSourceCheckConnectionJob(source, imageName, protocolVersion, isCustomConnector)); } public CheckConnectionRead checkSourceConnectionFromSourceCreate(final SourceCoreConfig sourceConfig) @@ -165,7 +166,8 @@ public CheckConnectionRead checkSourceConnectionFromSourceCreate(final SourceCor final Version protocolVersion = new Version(sourceDef.getProtocolVersion()); final String imageName = DockerUtils.getTaggedImageName(sourceDef.getDockerRepository(), sourceDef.getDockerImageTag()); - return reportConnectionStatus(synchronousSchedulerClient.createSourceCheckConnectionJob(source, imageName, protocolVersion)); + final boolean isCustomConnector = sourceDef.getCustom(); + return reportConnectionStatus(synchronousSchedulerClient.createSourceCheckConnectionJob(source, imageName, protocolVersion, isCustomConnector)); } public CheckConnectionRead checkSourceConnectionFromSourceIdForUpdate(final SourceUpdate sourceUpdate) @@ -188,8 +190,9 @@ public CheckConnectionRead checkDestinationConnectionFromDestinationId(final Des final DestinationConnection destination = configRepository.getDestinationConnection(destinationIdRequestBody.getDestinationId()); final StandardDestinationDefinition destinationDef = configRepository.getStandardDestinationDefinition(destination.getDestinationDefinitionId()); final String imageName = DockerUtils.getTaggedImageName(destinationDef.getDockerRepository(), destinationDef.getDockerImageTag()); + final boolean isCustomConnector = destinationDef.getCustom(); final Version protocolVersion = new Version(destinationDef.getProtocolVersion()); - return reportConnectionStatus(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion)); + return reportConnectionStatus(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion,isCustomConnector)); } public CheckConnectionRead checkDestinationConnectionFromDestinationCreate(final DestinationCoreConfig destinationConfig) @@ -198,6 +201,7 @@ public CheckConnectionRead checkDestinationConnectionFromDestinationCreate(final final var partialConfig = secretsRepositoryWriter.statefulSplitEphemeralSecrets( destinationConfig.getConnectionConfiguration(), destDef.getSpec()); + final boolean isCustomConnector = destDef.getCustom(); // todo (cgardens) - narrow the struct passed to the client. we are not setting fields that are // technically declared as required. @@ -205,10 +209,9 @@ public CheckConnectionRead checkDestinationConnectionFromDestinationCreate(final .withDestinationDefinitionId(destinationConfig.getDestinationDefinitionId()) .withConfiguration(partialConfig) .withWorkspaceId(destinationConfig.getWorkspaceId()); - final String imageName = DockerUtils.getTaggedImageName(destDef.getDockerRepository(), destDef.getDockerImageTag()); final Version protocolVersion = new Version(destDef.getProtocolVersion()); - return reportConnectionStatus(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion)); + return reportConnectionStatus(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion, isCustomConnector)); } public CheckConnectionRead checkDestinationConnectionFromDestinationIdForUpdate(final DestinationUpdate destinationUpdate) @@ -231,6 +234,7 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceId(final Source final SourceConnection source = configRepository.getSourceConnection(discoverSchemaRequestBody.getSourceId()); final StandardSourceDefinition sourceDef = configRepository.getStandardSourceDefinition(source.getSourceDefinitionId()); final String imageName = DockerUtils.getTaggedImageName(sourceDef.getDockerRepository(), sourceDef.getDockerImageTag()); + final boolean isCustomConnector = sourceDef.getCustom(); final String configHash = HASH_FUNCTION.hashBytes(Jsons.serialize(source.getConfiguration()).getBytes( Charsets.UTF_8)).toString(); @@ -240,7 +244,7 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceId(final Source final boolean bustActorCatalogCache = discoverSchemaRequestBody.getDisableCache() != null && discoverSchemaRequestBody.getDisableCache(); if (currentCatalog.isEmpty() || bustActorCatalogCache) { final SynchronousResponse persistedCatalogId = - synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, connectorVersion, new Version(sourceDef.getProtocolVersion())); + synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, connectorVersion, new Version(sourceDef.getProtocolVersion()), isCustomConnector); final SourceDiscoverSchemaRead discoveredSchema = retrieveDiscoveredSchema(persistedCatalogId); if (discoverSchemaRequestBody.getConnectionId() != null) { @@ -272,6 +276,7 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceCreate(final So sourceDef.getSpec()); final String imageName = DockerUtils.getTaggedImageName(sourceDef.getDockerRepository(), sourceDef.getDockerImageTag()); + final boolean isCustomConnector = sourceDef.getCustom(); // todo (cgardens) - narrow the struct passed to the client. we are not setting fields that are // technically declared as required. final SourceConnection source = new SourceConnection() @@ -279,7 +284,7 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceCreate(final So .withConfiguration(partialConfig) .withWorkspaceId(sourceCreate.getWorkspaceId()); final SynchronousResponse response = synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, sourceDef.getDockerImageTag(), - new Version(sourceDef.getProtocolVersion())); + new Version(sourceDef.getProtocolVersion()), isCustomConnector); return retrieveDiscoveredSchema(response); } diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java index 7d07cb59ec1a..e4702e7882b5 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java @@ -218,7 +218,7 @@ public SourceDefinitionRead createCustomSourceDefinition(final CustomSourceDefin private StandardSourceDefinition sourceDefinitionFromCreate(final SourceDefinitionCreate sourceDefinitionCreate) throws IOException { - final ConnectorSpecification spec = getSpecForImage(sourceDefinitionCreate.getDockerRepository(), sourceDefinitionCreate.getDockerImageTag()); + final ConnectorSpecification spec = getSpecForImage(sourceDefinitionCreate.getDockerRepository(), sourceDefinitionCreate.getDockerImageTag(), true); final Version airbyteProtocolVersion = AirbyteProtocolVersion.getWithDefault(spec.getProtocolVersion()); @@ -247,7 +247,7 @@ public SourceDefinitionRead updateSourceDefinition(final SourceDefinitionUpdate final boolean specNeedsUpdate = !currentSourceDefinition.getDockerImageTag().equals(sourceDefinitionUpdate.getDockerImageTag()) || sourceDefinitionUpdate.getDockerImageTag().equals(DEV_IMAGE_TAG); final ConnectorSpecification spec = specNeedsUpdate - ? getSpecForImage(currentSourceDefinition.getDockerRepository(), sourceDefinitionUpdate.getDockerImageTag()) + ? getSpecForImage(currentSourceDefinition.getDockerRepository(), sourceDefinitionUpdate.getDockerImageTag(), currentSourceDefinition.getCustom()) : currentSourceDefinition.getSpec(); final ActorDefinitionResourceRequirements updatedResourceReqs = sourceDefinitionUpdate.getResourceRequirements() != null ? ApiPojoConverters.actorDefResourceReqsToInternal(sourceDefinitionUpdate.getResourceRequirements()) @@ -317,9 +317,9 @@ public void deleteCustomSourceDefinition(final SourceDefinitionIdWithWorkspaceId deleteSourceDefinition(new SourceDefinitionIdRequestBody().sourceDefinitionId(definitionId)); } - private ConnectorSpecification getSpecForImage(final String dockerRepository, final String imageTag) throws IOException { + private ConnectorSpecification getSpecForImage(final String dockerRepository, final String imageTag, final boolean isCustomConnector) throws IOException { final String imageName = DockerUtils.getTaggedImageName(dockerRepository, imageTag); - final SynchronousResponse getSpecResponse = schedulerSynchronousClient.createGetSpecJob(imageName); + final SynchronousResponse getSpecResponse = schedulerSynchronousClient.createGetSpecJob(imageName, isCustomConnector); return SpecFetcher.getSpecFromJob(getSpecResponse); } diff --git a/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java b/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java index d96acb0dd5d2..0d250cf2c98b 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java +++ b/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java @@ -62,7 +62,8 @@ public DefaultSynchronousSchedulerClient(final TemporalClient temporalClient, @Override public SynchronousResponse createSourceCheckConnectionJob(final SourceConnection source, final String dockerImage, - final Version protocolVersion) + final Version protocolVersion, + final boolean isCustomConnector) throws IOException { final JsonNode sourceConfiguration = oAuthConfigSupplier.injectSourceOAuthParameters( source.getSourceDefinitionId(), @@ -71,7 +72,8 @@ public SynchronousResponse createSourceCheckConne final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig() .withConnectionConfiguration(sourceConfiguration) .withDockerImage(dockerImage) - .withProtocolVersion(protocolVersion); + .withProtocolVersion(protocolVersion) + .withIsCustomConnector(isCustomConnector); final UUID jobId = UUID.randomUUID(); final ConnectorJobReportingContext jobReportingContext = new ConnectorJobReportingContext(jobId, dockerImage); @@ -88,7 +90,8 @@ public SynchronousResponse createSourceCheckConne @Override public SynchronousResponse createDestinationCheckConnectionJob(final DestinationConnection destination, final String dockerImage, - final Version protocolVersion) + final Version protocolVersion, + final boolean isCustomConnector) throws IOException { final JsonNode destinationConfiguration = oAuthConfigSupplier.injectDestinationOAuthParameters( destination.getDestinationDefinitionId(), @@ -97,7 +100,8 @@ public SynchronousResponse createDestinationCheck final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig() .withConnectionConfiguration(destinationConfiguration) .withDockerImage(dockerImage) - .withProtocolVersion(protocolVersion); + .withProtocolVersion(protocolVersion) + .withIsCustomConnector(isCustomConnector); final UUID jobId = UUID.randomUUID(); final ConnectorJobReportingContext jobReportingContext = new ConnectorJobReportingContext(jobId, dockerImage); @@ -115,7 +119,8 @@ public SynchronousResponse createDestinationCheck public SynchronousResponse createDiscoverSchemaJob(final SourceConnection source, final String dockerImage, final String connectorVersion, - final Version protocolVersion) + final Version protocolVersion, + final boolean isCustomConnector) throws IOException { final JsonNode sourceConfiguration = oAuthConfigSupplier.injectSourceOAuthParameters( source.getSourceDefinitionId(), @@ -128,7 +133,8 @@ public SynchronousResponse createDiscoverSchemaJob(final SourceConnection .withSourceId(source.getSourceId().toString()) .withConfigHash(HASH_FUNCTION.hashBytes(Jsons.serialize(source.getConfiguration()).getBytes( Charsets.UTF_8)).toString()) - .withConnectorVersion(connectorVersion); + .withConnectorVersion(connectorVersion) + .withIsCustomConnector(isCustomConnector); final UUID jobId = UUID.randomUUID(); final ConnectorJobReportingContext jobReportingContext = new ConnectorJobReportingContext(jobId, dockerImage); @@ -143,8 +149,8 @@ public SynchronousResponse createDiscoverSchemaJob(final SourceConnection } @Override - public SynchronousResponse createGetSpecJob(final String dockerImage) throws IOException { - final JobGetSpecConfig jobSpecConfig = new JobGetSpecConfig().withDockerImage(dockerImage); + public SynchronousResponse createGetSpecJob(final String dockerImage, final boolean isCustomConnector) throws IOException { + final JobGetSpecConfig jobSpecConfig = new JobGetSpecConfig().withDockerImage(dockerImage).withIsCustomConnector(isCustomConnector); final UUID jobId = UUID.randomUUID(); final ConnectorJobReportingContext jobReportingContext = new ConnectorJobReportingContext(jobId, dockerImage); diff --git a/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java b/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java index f472d9da7ced..e794fbf54579 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java +++ b/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java @@ -20,17 +20,17 @@ public interface SynchronousSchedulerClient { SynchronousResponse createSourceCheckConnectionJob(SourceConnection source, String dockerImage, - Version protocolVersion) + Version protocolVersion,boolean isCustomConnector) throws IOException; SynchronousResponse createDestinationCheckConnectionJob(DestinationConnection destination, String dockerImage, - Version protocolVersion) + Version protocolVersion,boolean isCustomConnector) throws IOException; - SynchronousResponse createDiscoverSchemaJob(SourceConnection source, String dockerImage, String connectorVersion, Version protocolVersion) + SynchronousResponse createDiscoverSchemaJob(SourceConnection source, String dockerImage, String connectorVersion, Version protocolVersion, boolean isCustomConnector) throws IOException; - SynchronousResponse createGetSpecJob(String dockerImage) throws IOException; + SynchronousResponse createGetSpecJob(String dockerImage,boolean isCustomConnector) throws IOException; } diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/DestinationDefinitionsHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/DestinationDefinitionsHandlerTest.java index daad65516f51..9477b8a028f7 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/DestinationDefinitionsHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/DestinationDefinitionsHandlerTest.java @@ -343,7 +343,7 @@ void testCreateDestinationDefinition() throws URISyntaxException, IOException, J final String imageName = DockerUtils.getTaggedImageName(destination.getDockerRepository(), destination.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(destination.getDestinationDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( destination.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -375,7 +375,7 @@ void testCreateDestinationDefinition() throws URISyntaxException, IOException, J final DestinationDefinitionRead actualRead = destinationDefinitionsHandler.createPrivateDestinationDefinition(create); assertEquals(expectedRead, actualRead); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository).writeStandardDestinationDefinition(destination .withProtocolVersion(DEFAULT_PROTOCOL_VERSION) .withReleaseDate(null) @@ -391,7 +391,7 @@ void testCreateDestinationDefinitionShouldCheckProtocolVersion() throws URISynta final String imageName = DockerUtils.getTaggedImageName(destination.getDockerRepository(), destination.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(destination.getDestinationDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( destination.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -408,7 +408,7 @@ void testCreateDestinationDefinitionShouldCheckProtocolVersion() throws URISynta assertThrows(UnsupportedProtocolVersionException.class, () -> destinationDefinitionsHandler.createPrivateDestinationDefinition(create)); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository, never()).writeStandardDestinationDefinition(destination .withProtocolVersion(DEFAULT_PROTOCOL_VERSION) .withReleaseDate(null) @@ -422,7 +422,7 @@ void testCreateCustomDestinationDefinition() throws URISyntaxException, IOExcept final String imageName = DockerUtils.getTaggedImageName(destination.getDockerRepository(), destination.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(destination.getDestinationDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( destination.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -458,7 +458,7 @@ void testCreateCustomDestinationDefinition() throws URISyntaxException, IOExcept final DestinationDefinitionRead actualRead = destinationDefinitionsHandler.createCustomDestinationDefinition(customCreate); assertEquals(expectedRead, actualRead); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository).writeCustomDestinationDefinition( destination .withProtocolVersion(DEFAULT_PROTOCOL_VERSION) @@ -477,7 +477,7 @@ void testCreateCustomDestinationDefinitionWithInvalidProtocol() throws URISyntax final String imageName = DockerUtils.getTaggedImageName(destination.getDockerRepository(), destination.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(destination.getDestinationDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( destination.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -498,7 +498,7 @@ void testCreateCustomDestinationDefinitionWithInvalidProtocol() throws URISyntax assertThrows(UnsupportedProtocolVersionException.class, () -> destinationDefinitionsHandler.createCustomDestinationDefinition(customCreate)); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository, never()).writeCustomDestinationDefinition( destination .withProtocolVersion(invalidProtocol) @@ -525,7 +525,7 @@ void testUpdateDestination() throws ConfigNotFoundException, IOException, JsonVa final ConnectorSpecification newSpec = new ConnectorSpecification() .withConnectionSpecification(Jsons.jsonNode(ImmutableMap.of("foo2", "bar2"))) .withProtocolVersion(newProtocolVersion); - when(schedulerSynchronousClient.createGetSpecJob(newImageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(newImageName, false)).thenReturn(new SynchronousResponse<>( newSpec, SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -537,7 +537,7 @@ void testUpdateDestination() throws ConfigNotFoundException, IOException, JsonVa .dockerImageTag(newDockerImageTag)); assertEquals(newDockerImageTag, destinationRead.getDockerImageTag()); - verify(schedulerSynchronousClient).createGetSpecJob(newImageName); + verify(schedulerSynchronousClient).createGetSpecJob(newImageName, false); verify(configRepository).writeStandardDestinationDefinition(updatedDestination); final Configs configs = new EnvConfigs(); @@ -564,7 +564,7 @@ void testOutOfProtocolRangeUpdateDestination() throws ConfigNotFoundException, I final ConnectorSpecification newSpec = new ConnectorSpecification() .withConnectionSpecification(Jsons.jsonNode(ImmutableMap.of("foo2", "bar2"))) .withProtocolVersion(newProtocolVersion); - when(schedulerSynchronousClient.createGetSpecJob(newImageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(newImageName, false)).thenReturn(new SynchronousResponse<>( newSpec, SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -575,7 +575,7 @@ void testOutOfProtocolRangeUpdateDestination() throws ConfigNotFoundException, I new DestinationDefinitionUpdate().destinationDefinitionId(this.destinationDefinition.getDestinationDefinitionId()) .dockerImageTag(newDockerImageTag))); - verify(schedulerSynchronousClient).createGetSpecJob(newImageName); + verify(schedulerSynchronousClient).createGetSpecJob(newImageName, false); verify(configRepository, never()).writeStandardDestinationDefinition(updatedDestination); } diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java index 2eb3e9be56d7..bbe304967dc4 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java @@ -190,13 +190,13 @@ void testCheckSourceConnectionFromSourceId() throws JsonValidationException, IOE .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion)) + when(synchronousSchedulerClient.createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion, false)) .thenReturn((SynchronousResponse) jobResponse); schedulerHandler.checkSourceConnectionFromSourceId(request); verify(configRepository).getSourceConnection(source.getSourceId()); - verify(synchronousSchedulerClient).createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion); + verify(synchronousSchedulerClient).createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion, false); } @Test @@ -221,12 +221,12 @@ void testCheckSourceConnectionFromSourceCreate() throws JsonValidationException, when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(source.getConfiguration()), any())).thenReturn(source.getConfiguration()); - when(synchronousSchedulerClient.createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion)) + when(synchronousSchedulerClient.createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion, false)) .thenReturn((SynchronousResponse) jobResponse); schedulerHandler.checkSourceConnectionFromSourceCreate(sourceCoreConfig); - verify(synchronousSchedulerClient).createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion); + verify(synchronousSchedulerClient).createSourceCheckConnectionJob(source, SOURCE_DOCKER_IMAGE, protocolVersion, false); } @Test @@ -250,7 +250,7 @@ void testCheckSourceConnectionFromUpdate() throws IOException, JsonValidationExc final SourceConnection submittedSource = new SourceConnection() .withSourceDefinitionId(source.getSourceDefinitionId()) .withConfiguration(source.getConfiguration()); - when(synchronousSchedulerClient.createSourceCheckConnectionJob(submittedSource, DESTINATION_DOCKER_IMAGE, protocolVersion)) + when(synchronousSchedulerClient.createSourceCheckConnectionJob(submittedSource, DESTINATION_DOCKER_IMAGE, protocolVersion, false)) .thenReturn((SynchronousResponse) jobResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(source.getConfiguration()), @@ -258,7 +258,7 @@ void testCheckSourceConnectionFromUpdate() throws IOException, JsonValidationExc schedulerHandler.checkSourceConnectionFromSourceIdForUpdate(sourceUpdate); verify(jsonSchemaValidator).ensure(CONNECTOR_SPECIFICATION.getConnectionSpecification(), source.getConfiguration()); - verify(synchronousSchedulerClient).createSourceCheckConnectionJob(submittedSource, DESTINATION_DOCKER_IMAGE, protocolVersion); + verify(synchronousSchedulerClient).createSourceCheckConnectionJob(submittedSource, DESTINATION_DOCKER_IMAGE, protocolVersion, false); } @Test @@ -314,14 +314,14 @@ void testCheckDestinationConnectionFromDestinationId() throws IOException, JsonV .withDestinationDefinitionId(destination.getDestinationDefinitionId())); when(configRepository.getDestinationConnection(destination.getDestinationId())).thenReturn(destination); when(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, DESTINATION_DOCKER_IMAGE, - new Version(DESTINATION_PROTOCOL_VERSION))) + new Version(DESTINATION_PROTOCOL_VERSION), false)) .thenReturn((SynchronousResponse) jobResponse); schedulerHandler.checkDestinationConnectionFromDestinationId(request); verify(configRepository).getDestinationConnection(destination.getDestinationId()); verify(synchronousSchedulerClient).createDestinationCheckConnectionJob(destination, DESTINATION_DOCKER_IMAGE, - new Version(DESTINATION_PROTOCOL_VERSION)); + new Version(DESTINATION_PROTOCOL_VERSION), false); } @Test @@ -342,7 +342,7 @@ void testCheckDestinationConnectionFromDestinationCreate() throws JsonValidation .withDestinationDefinitionId(destination.getDestinationDefinitionId())); when(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, DESTINATION_DOCKER_IMAGE, - new Version(DESTINATION_PROTOCOL_VERSION))) + new Version(DESTINATION_PROTOCOL_VERSION), false)) .thenReturn((SynchronousResponse) jobResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(destination.getConfiguration()), @@ -350,7 +350,7 @@ void testCheckDestinationConnectionFromDestinationCreate() throws JsonValidation schedulerHandler.checkDestinationConnectionFromDestinationCreate(destinationCoreConfig); verify(synchronousSchedulerClient).createDestinationCheckConnectionJob(destination, DESTINATION_DOCKER_IMAGE, - new Version(DESTINATION_PROTOCOL_VERSION)); + new Version(DESTINATION_PROTOCOL_VERSION), false); } @Test @@ -375,7 +375,7 @@ void testCheckDestinationConnectionFromUpdate() throws IOException, JsonValidati .withDestinationDefinitionId(destination.getDestinationDefinitionId()) .withConfiguration(destination.getConfiguration()); when(synchronousSchedulerClient.createDestinationCheckConnectionJob(submittedDestination, DESTINATION_DOCKER_IMAGE, - new Version(DESTINATION_PROTOCOL_VERSION))) + new Version(DESTINATION_PROTOCOL_VERSION), false)) .thenReturn((SynchronousResponse) jobResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(destination.getConfiguration()), @@ -384,7 +384,7 @@ void testCheckDestinationConnectionFromUpdate() throws IOException, JsonValidati verify(jsonSchemaValidator).ensure(CONNECTOR_SPECIFICATION.getConnectionSpecification(), destination.getConfiguration()); verify(synchronousSchedulerClient).createDestinationCheckConnectionJob(submittedDestination, DESTINATION_DOCKER_IMAGE, - new Version(DESTINATION_PROTOCOL_VERSION)); + new Version(DESTINATION_PROTOCOL_VERSION), false); } @Test @@ -412,7 +412,7 @@ void testDiscoverSchemaForSourceFromSourceId() throws IOException, JsonValidatio .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); when(configRepository.getActorCatalog(any(), any(), any())).thenReturn(Optional.empty()); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn(discoverResponse); final SourceDiscoverSchemaRead actual = schedulerHandler.discoverSchemaForSourceFromSourceId(request); @@ -423,7 +423,7 @@ void testDiscoverSchemaForSourceFromSourceId() throws IOException, JsonValidatio assertTrue(actual.getJobInfo().getSucceeded()); verify(configRepository).getSourceConnection(source.getSourceId()); verify(configRepository).getActorCatalog(eq(request.getSourceId()), eq(SOURCE_DOCKER_TAG), any()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION)); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); } @Test @@ -450,7 +450,7 @@ void testDiscoverSchemaForSourceFromSourceIdCachedCatalog() throws IOException, .withCatalogHash("") .withId(thisCatalogId); when(configRepository.getActorCatalog(any(), any(), any())).thenReturn(Optional.of(actorCatalog)); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn(discoverResponse); final SourceDiscoverSchemaRead actual = schedulerHandler.discoverSchemaForSourceFromSourceId(request); @@ -463,7 +463,7 @@ void testDiscoverSchemaForSourceFromSourceIdCachedCatalog() throws IOException, verify(configRepository).getActorCatalog(eq(request.getSourceId()), any(), any()); verify(configRepository, never()).writeActorCatalogFetchEvent(any(), any(), any(), any()); verify(synchronousSchedulerClient, never()).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, - new Version(SOURCE_PROTOCOL_VERSION)); + new Version(SOURCE_PROTOCOL_VERSION), false); } @Test @@ -491,7 +491,7 @@ void testDiscoverSchemaForSourceFromSourceIdDisableCache() throws IOException, J .withCatalogHash("") .withId(discoveredCatalogId); when(configRepository.getActorCatalogById(discoveredCatalogId)).thenReturn(actorCatalog); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn(discoverResponse); final SourceDiscoverSchemaRead actual = schedulerHandler.discoverSchemaForSourceFromSourceId(request); @@ -501,7 +501,7 @@ void testDiscoverSchemaForSourceFromSourceIdDisableCache() throws IOException, J assertTrue(actual.getJobInfo().getSucceeded()); verify(configRepository).getSourceConnection(source.getSourceId()); verify(configRepository).getActorCatalog(eq(request.getSourceId()), any(), any()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION)); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); } @Test @@ -516,7 +516,7 @@ void testDiscoverSchemaForSourceFromSourceIdFailed() throws IOException, JsonVal .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn((SynchronousResponse) jobResponse); when(completedJob.getSuccessOutput()).thenReturn(Optional.empty()); when(completedJob.getStatus()).thenReturn(JobStatus.FAILED); @@ -527,7 +527,7 @@ void testDiscoverSchemaForSourceFromSourceIdFailed() throws IOException, JsonVal assertNotNull(actual.getJobInfo()); assertFalse(actual.getJobInfo().getSucceeded()); verify(configRepository).getSourceConnection(source.getSourceId()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION)); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); } @Test @@ -548,7 +548,7 @@ void testDiscoverSchemaFromSourceIdWithConnectionIdNonBreaking() throws IOExcept .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn(discoverResponse); when(discoverResponse.isSuccess()).thenReturn(true); @@ -596,7 +596,7 @@ void testDiscoverSchemaFromSourceIdWithConnectionIdBreaking() throws IOException .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn(discoverResponse); when(discoverResponse.isSuccess()).thenReturn(true); @@ -655,7 +655,7 @@ void testDiscoverSchemaForSourceFromSourceCreate() throws JsonValidationExceptio .withDockerImageTag(SOURCE_DOCKER_TAG) .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn(discoverResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(source.getConfiguration()), @@ -667,7 +667,7 @@ void testDiscoverSchemaForSourceFromSourceCreate() throws JsonValidationExceptio assertNotNull(actual.getJobInfo()); assertEquals(actual.getCatalogId(), discoverResponse.getOutput()); assertTrue(actual.getJobInfo().getSucceeded()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION)); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); } @Test @@ -687,7 +687,7 @@ void testDiscoverSchemaForSourceFromSourceCreateFailed() throws JsonValidationEx .withDockerImageTag(SOURCE_DOCKER_TAG) .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION))) + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) .thenReturn((SynchronousResponse) jobResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(source.getConfiguration()), @@ -700,7 +700,7 @@ void testDiscoverSchemaForSourceFromSourceCreateFailed() throws JsonValidationEx assertNull(actual.getCatalog()); assertNotNull(actual.getJobInfo()); assertFalse(actual.getJobInfo().getSucceeded()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION)); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); } @Test diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/SourceDefinitionsHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/SourceDefinitionsHandlerTest.java index 733c50995ca7..d3c684476618 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/SourceDefinitionsHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/SourceDefinitionsHandlerTest.java @@ -330,7 +330,7 @@ void testCreateSourceDefinition() throws URISyntaxException, IOException, JsonVa final String imageName = DockerUtils.getTaggedImageName(sourceDefinition.getDockerRepository(), sourceDefinition.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(sourceDefinition.getSourceDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( sourceDefinition.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -362,7 +362,7 @@ void testCreateSourceDefinition() throws URISyntaxException, IOException, JsonVa final SourceDefinitionRead actualRead = sourceDefinitionsHandler.createPrivateSourceDefinition(create); assertEquals(expectedRead, actualRead); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository) .writeStandardSourceDefinition( sourceDefinition @@ -380,7 +380,7 @@ void testCreateSourceDefinitionWithInvalidProtocol() throws URISyntaxException, final String imageName = DockerUtils.getTaggedImageName(sourceDefinition.getDockerRepository(), sourceDefinition.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(sourceDefinition.getSourceDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( sourceDefinition.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -397,7 +397,7 @@ void testCreateSourceDefinitionWithInvalidProtocol() throws URISyntaxException, assertThrows(UnsupportedProtocolVersionException.class, () -> sourceDefinitionsHandler.createPrivateSourceDefinition(create)); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository, never()) .writeStandardSourceDefinition( sourceDefinition @@ -413,7 +413,7 @@ void testCreateCustomSourceDefinition() throws URISyntaxException, IOException, final String imageName = DockerUtils.getTaggedImageName(sourceDefinition.getDockerRepository(), sourceDefinition.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(sourceDefinition.getSourceDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( sourceDefinition.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -449,7 +449,7 @@ void testCreateCustomSourceDefinition() throws URISyntaxException, IOException, final SourceDefinitionRead actualRead = sourceDefinitionsHandler.createCustomSourceDefinition(customCreate); assertEquals(expectedRead, actualRead); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository).writeCustomSourceDefinition( sourceDefinition .withReleaseDate(null) @@ -468,7 +468,7 @@ void testCreateCustomSourceDefinitionWithInvalidProtocol() throws URISyntaxExcep final String imageName = DockerUtils.getTaggedImageName(sourceDefinition.getDockerRepository(), sourceDefinition.getDockerImageTag()); when(uuidSupplier.get()).thenReturn(sourceDefinition.getSourceDefinitionId()); - when(schedulerSynchronousClient.createGetSpecJob(imageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(imageName, true)).thenReturn(new SynchronousResponse<>( sourceDefinition.getSpec(), SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -489,7 +489,7 @@ void testCreateCustomSourceDefinitionWithInvalidProtocol() throws URISyntaxExcep assertThrows(UnsupportedProtocolVersionException.class, () -> sourceDefinitionsHandler.createCustomSourceDefinition(customCreate)); - verify(schedulerSynchronousClient).createGetSpecJob(imageName); + verify(schedulerSynchronousClient).createGetSpecJob(imageName, true); verify(configRepository, never()).writeCustomSourceDefinition( sourceDefinition .withReleaseDate(null) @@ -514,7 +514,7 @@ void testUpdateSourceDefinition() throws ConfigNotFoundException, IOException, J final ConnectorSpecification newSpec = new ConnectorSpecification() .withConnectionSpecification(Jsons.jsonNode(ImmutableMap.of("foo2", "bar2"))) .withProtocolVersion(newProtocolVersion); - when(schedulerSynchronousClient.createGetSpecJob(newImageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(newImageName, false)).thenReturn(new SynchronousResponse<>( newSpec, SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -526,7 +526,7 @@ void testUpdateSourceDefinition() throws ConfigNotFoundException, IOException, J new SourceDefinitionUpdate().sourceDefinitionId(this.sourceDefinition.getSourceDefinitionId()).dockerImageTag(newDockerImageTag)); assertEquals(newDockerImageTag, sourceDefinitionRead.getDockerImageTag()); - verify(schedulerSynchronousClient).createGetSpecJob(newImageName); + verify(schedulerSynchronousClient).createGetSpecJob(newImageName, false); verify(configRepository).writeStandardSourceDefinition(updatedSource); final Configs configs = new EnvConfigs(); @@ -550,7 +550,7 @@ void testUpdateSourceDefinitionWithInvalidProtocol() throws ConfigNotFoundExcept final ConnectorSpecification newSpec = new ConnectorSpecification() .withConnectionSpecification(Jsons.jsonNode(ImmutableMap.of("foo2", "bar2"))) .withProtocolVersion(newProtocolVersion); - when(schedulerSynchronousClient.createGetSpecJob(newImageName)).thenReturn(new SynchronousResponse<>( + when(schedulerSynchronousClient.createGetSpecJob(newImageName, false)).thenReturn(new SynchronousResponse<>( newSpec, SynchronousJobMetadata.mock(ConfigType.GET_SPEC))); @@ -561,7 +561,7 @@ void testUpdateSourceDefinitionWithInvalidProtocol() throws ConfigNotFoundExcept .updateSourceDefinition( new SourceDefinitionUpdate().sourceDefinitionId(this.sourceDefinition.getSourceDefinitionId()).dockerImageTag(newDockerImageTag))); - verify(schedulerSynchronousClient).createGetSpecJob(newImageName); + verify(schedulerSynchronousClient).createGetSpecJob(newImageName, false); verify(configRepository, never()).writeStandardSourceDefinition(updatedSource); } diff --git a/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java b/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java index e9af5118545d..25586aa1b89c 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java @@ -204,14 +204,15 @@ void testCreateSourceCheckConnectionJob() throws IOException { final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig() .withConnectionConfiguration(SOURCE_CONNECTION.getConfiguration()) .withDockerImage(DOCKER_IMAGE) - .withProtocolVersion(PROTOCOL_VERSION); + .withProtocolVersion(PROTOCOL_VERSION) .withIsCustomConnector(false); + final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class); final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput); when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(jobCheckConnectionConfig))) .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true))); final SynchronousResponse response = - schedulerClient.createSourceCheckConnectionJob(SOURCE_CONNECTION, DOCKER_IMAGE, PROTOCOL_VERSION); + schedulerClient.createSourceCheckConnectionJob(SOURCE_CONNECTION, DOCKER_IMAGE, PROTOCOL_VERSION, false); assertEquals(mockOutput, response.getOutput()); } @@ -220,14 +221,15 @@ void testCreateDestinationCheckConnectionJob() throws IOException { final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig() .withConnectionConfiguration(DESTINATION_CONNECTION.getConfiguration()) .withDockerImage(DOCKER_IMAGE) - .withProtocolVersion(PROTOCOL_VERSION); + .withProtocolVersion(PROTOCOL_VERSION) + .withIsCustomConnector(false); final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class); final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput); when(temporalClient.submitCheckConnection(any(UUID.class), eq(0), eq(jobCheckConnectionConfig))) .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true))); final SynchronousResponse response = - schedulerClient.createDestinationCheckConnectionJob(DESTINATION_CONNECTION, DOCKER_IMAGE, PROTOCOL_VERSION); + schedulerClient.createDestinationCheckConnectionJob(DESTINATION_CONNECTION, DOCKER_IMAGE, PROTOCOL_VERSION, false); assertEquals(mockOutput, response.getOutput()); } @@ -238,19 +240,19 @@ void testCreateDiscoverSchemaJob() throws IOException { when(temporalClient.submitDiscoverSchema(any(UUID.class), eq(0), any(JobDiscoverCatalogConfig.class))) .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true))); final SynchronousResponse response = - schedulerClient.createDiscoverSchemaJob(SOURCE_CONNECTION, DOCKER_IMAGE, DOCKER_IMAGE_TAG, PROTOCOL_VERSION); + schedulerClient.createDiscoverSchemaJob(SOURCE_CONNECTION, DOCKER_IMAGE, DOCKER_IMAGE_TAG, PROTOCOL_VERSION, false); assertEquals(expectedCatalogId, response.getOutput()); } @Test void testCreateGetSpecJob() throws IOException { - final JobGetSpecConfig jobSpecConfig = new JobGetSpecConfig().withDockerImage(DOCKER_IMAGE); + final JobGetSpecConfig jobSpecConfig = new JobGetSpecConfig().withDockerImage(DOCKER_IMAGE).withIsCustomConnector(false); final ConnectorSpecification mockOutput = mock(ConnectorSpecification.class); final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withSpec(mockOutput); when(temporalClient.submitGetSpec(any(UUID.class), eq(0), eq(jobSpecConfig))) .thenReturn(new TemporalResponse<>(jobOutput, createMetadata(true))); - final SynchronousResponse response = schedulerClient.createGetSpecJob(DOCKER_IMAGE); + final SynchronousResponse response = schedulerClient.createGetSpecJob(DOCKER_IMAGE, false); assertEquals(mockOutput, response.getOutput()); } diff --git a/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml b/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml index a4c7b6fa63c6..5dcf6d11af0c 100644 --- a/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml +++ b/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml @@ -19,3 +19,5 @@ properties: protocolVersion: type: object existingJavaType: io.airbyte.commons.version.Version + isCustomConnector: + type: boolean \ No newline at end of file diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index fbee2323e0f9..07827acf7180 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -206,6 +206,7 @@ public WorkerConfigs checkWorkerConfigs( @Named("checkResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("checkJobKubeNodeSelectors") final Map nodeSelectors, + @Named("customNodeSelectors") final Map customNodeSelectors, @Named("checkJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -219,6 +220,7 @@ public WorkerConfigs checkWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, + customNodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -236,7 +238,9 @@ public WorkerConfigs defaultWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, - @Named("defaultJobKubeAnnotations") final Map annotations, + @Named("customNodeSelectors") final Map customNodeSelectors, + + @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -249,6 +253,7 @@ public WorkerConfigs defaultWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, + customNodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -267,7 +272,9 @@ public WorkerConfigs discoverWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("discoverJobKubeNodeSelectors") final Map nodeSelectors, - @Named("discoverJobKubeAnnotations") final Map annotations, + @Named("customNodeSelectors") final Map customNodeSelectors, + + @Named("discoverJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -280,6 +287,7 @@ public WorkerConfigs discoverWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, + customNodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -297,7 +305,9 @@ public WorkerConfigs replicationWorkerConfigs( @Named("replicationResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, - @Named("defaultJobKubeAnnotations") final Map annotations, + @Named("customNodeSelectors") final Map customNodeSelectors, + + @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -310,6 +320,7 @@ public WorkerConfigs replicationWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, + customNodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -328,7 +339,9 @@ public WorkerConfigs specWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("specJobKubeNodeSelectors") final Map nodeSelectors, - @Named("specJobKubeAnnotations") final Map annotations, + @Named("customNodeSelectors") final Map customNodeSelectors, + + @Named("specJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -341,6 +354,7 @@ public WorkerConfigs specWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, + customNodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java index 167ae3ac7282..ba21f2af375a 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java @@ -130,7 +130,8 @@ private CheckedSupplier Math.toIntExact(launcherConfig.getAttemptId()), launcherConfig.getDockerImage(), processFactory, - workerConfigs.getResourceRequirements()); + workerConfigs.getResourceRequirements(), + launcherConfig.getIsCustomConnector()); final AirbyteStreamFactory streamFactory = launcherConfig.getProtocolVersion() != null ? new VersionedAirbyteStreamFactory<>(serDeProvider, migratorFactory, launcherConfig.getProtocolVersion()) : new DefaultAirbyteStreamFactory(); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java index 9ba130aed131..5a36c144fc43 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java @@ -124,7 +124,7 @@ private CheckedSupplier return () -> { final IntegrationLauncher integrationLauncher = new AirbyteIntegrationLauncher(launcherConfig.getJobId(), launcherConfig.getAttemptId().intValue(), launcherConfig.getDockerImage(), - processFactory, workerConfigs.getResourceRequirements()); + processFactory, workerConfigs.getResourceRequirements(), launcherConfig.getIsCustomConnector()); final AirbyteStreamFactory streamFactory = new VersionedAirbyteStreamFactory<>(serDeProvider, migratorFactory, launcherConfig.getProtocolVersion()); return new DefaultDiscoverCatalogWorker(configRepository, integrationLauncher, streamFactory); 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 index de7c26da10f5..89103b866f81 100644 --- 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 @@ -69,7 +69,9 @@ public GeneratedJobInput getSyncWorkflowInput(final SyncInput input) { .withConfiguredAirbyteCatalog(resetConnection.getConfiguredAirbyteCatalog()) .withOperationSequence(resetConnection.getOperationSequence()) .withResourceRequirements(resetConnection.getResourceRequirements()) - .withState(resetConnection.getState()); + .withState(resetConnection.getState()) + .withIsSourceCustomConnector(resetConnection.getIsSourceCustomConnector()) + .withIsDestinationCustomConnector(resetConnection.getIsDestinationCustomConnector()); } else { throw new IllegalStateException( String.format("Unexpected config type %s for job %d. The only supported config types for this activity are (%s)", @@ -84,13 +86,15 @@ public GeneratedJobInput getSyncWorkflowInput(final SyncInput input) { .withJobId(String.valueOf(jobId)) .withAttemptId((long) attempt) .withDockerImage(config.getSourceDockerImage()) - .withProtocolVersion(config.getSourceProtocolVersion()); + .withProtocolVersion(config.getSourceProtocolVersion()) + .withIsCustomConnector(config.getIsSourceCustomConnector()); final IntegrationLauncherConfig destinationLauncherConfig = new IntegrationLauncherConfig() .withJobId(String.valueOf(jobId)) .withAttemptId((long) attempt) .withDockerImage(config.getDestinationDockerImage()) - .withProtocolVersion(config.getDestinationProtocolVersion()); + .withProtocolVersion(config.getDestinationProtocolVersion()) + .withIsCustomConnector(config.getIsDestinationCustomConnector()); final StandardSyncInput syncInput = new StandardSyncInput() .withNamespaceDefinition(config.getNamespaceDefinition()) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java index 340eafae4d45..b90cf33c2347 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java @@ -83,7 +83,7 @@ public ConnectorJobOutput run(final JobRunConfig jobRunConfig, final Integration ApmTraceUtils.addTagsToTrace(Map.of(ATTEMPT_NUMBER_KEY, jobRunConfig.getAttemptId(), DOCKER_IMAGE_KEY, launcherConfig.getDockerImage(), JOB_ID_KEY, jobRunConfig.getJobId())); - final Supplier inputSupplier = () -> new JobGetSpecConfig().withDockerImage(launcherConfig.getDockerImage()); + final Supplier inputSupplier = () -> new JobGetSpecConfig().withDockerImage(launcherConfig.getDockerImage()).withIsCustomConnector(launcherConfig.getIsCustomConnector()); final ActivityExecutionContext context = Activity.getExecutionContext(); @@ -111,7 +111,8 @@ private CheckedSupplier, Exception> launcherConfig.getAttemptId().intValue(), launcherConfig.getDockerImage(), processFactory, - workerConfigs.getResourceRequirements()); + workerConfigs.getResourceRequirements(), + launcherConfig.getIsCustomConnector()); return new DefaultGetSpecWorker(integrationLauncher, streamFactory); }; 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 index b7104497e71a..ec438804b687 100644 --- 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 @@ -263,13 +263,15 @@ private CheckedSupplier, Exception> Math.toIntExact(sourceLauncherConfig.getAttemptId()), sourceLauncherConfig.getDockerImage(), processFactory, - syncInput.getSourceResourceRequirements()); + syncInput.getSourceResourceRequirements(), + sourceLauncherConfig.getIsCustomConnector()); final IntegrationLauncher destinationLauncher = new AirbyteIntegrationLauncher( destinationLauncherConfig.getJobId(), Math.toIntExact(destinationLauncherConfig.getAttemptId()), destinationLauncherConfig.getDockerImage(), processFactory, - syncInput.getDestinationResourceRequirements()); + syncInput.getDestinationResourceRequirements(), + destinationLauncherConfig.getIsCustomConnector()); // reset jobs use an empty source to induce resetting all data in destination. final AirbyteSource airbyteSource = diff --git a/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java b/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java index 7dcbdb3cdfb5..582004aef227 100644 --- a/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java +++ b/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java @@ -430,7 +430,7 @@ private Process getProcess(final String entrypoint, final Map fi private Process getProcess(final Map customLabels, final String entrypoint, final Map files) throws WorkerException { - return processFactory.create("tester", "some-id", 0, Path.of("/tmp/job-root"), "busybox:latest", false, files, entrypoint, + return processFactory.create("tester", "some-id", 0, Path.of("/tmp/job-root"), "busybox:latest", false, false, files, entrypoint, DEFAULT_RESOURCE_REQUIREMENTS, customLabels, Collections.emptyMap(), Collections.emptyMap()); } From 6f5ac3fda0cfb714ff763d2ecb05a8021792ec99 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 22 Nov 2022 13:46:00 -0800 Subject: [PATCH 03/18] config, style, pmd --- .../io/airbyte/workers/WorkerConfigs.java | 4 +- .../general/DbtTransformationRunner.java | 3 +- .../process/AirbyteIntegrationLauncher.java | 5 +- .../workers/process/KubeProcessFactory.java | 2 - .../AirbyteIntegrationLauncherTest.java | 2 +- .../types/JobCheckConnectionConfig.yaml | 2 +- .../resources/types/JobGetSpecConfig.yaml | 2 +- .../types/JobResetConnectionConfig.yaml | 2 +- .../main/resources/types/JobSyncConfig.yaml | 2 +- .../server/handlers/SchedulerHandler.java | 9 ++- .../handlers/SourceDefinitionsHandler.java | 9 ++- .../DefaultSynchronousSchedulerClient.java | 6 +- .../scheduler/SynchronousSchedulerClient.java | 14 +++-- .../server/handlers/SchedulerHandlerTest.java | 55 ++++++++++++------- ...DefaultSynchronousSchedulerClientTest.java | 3 +- .../IntegrationLauncherConfig.yaml | 2 +- .../WorkerConfigurationBeanFactory.java | 24 +++++--- .../temporal/spec/SpecActivityImpl.java | 3 +- .../src/main/resources/application.yml | 3 + 19 files changed, 93 insertions(+), 59 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java index 71dee3125214..de3355920a02 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java @@ -186,7 +186,9 @@ public Map getworkerKubeNodeSelectors() { return workerKubeNodeSelectors; } - public Map getWorkerIsolatedKubeNodeSelectors() {return workerIsolatedKubeNodeSelectors;} + public Map getWorkerIsolatedKubeNodeSelectors() { + return workerIsolatedKubeNodeSelectors; + } public Map getWorkerKubeAnnotations() { return workerKubeAnnotations; diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java index 708ce9f7a0a3..62316d0df602 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/general/DbtTransformationRunner.java @@ -98,7 +98,8 @@ public boolean transform(final String jobId, attempt, jobRoot, dbtConfig.getDockerImage(), - false, false, + false, + false, files, "/bin/bash", resourceRequirements, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java index 6cf05cf952fa..7844cafc0ac9 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java @@ -57,14 +57,14 @@ public class AirbyteIntegrationLauncher implements IntegrationLauncher { private final ResourceRequirements resourceRequirement; private final FeatureFlags featureFlags; - private boolean useIsolatedNodePool; + private final boolean useIsolatedNodePool; public AirbyteIntegrationLauncher(final String jobId, final int attempt, final String imageName, final ProcessFactory processFactory, final ResourceRequirements resourceRequirement, - final boolean useIsolatedNodePool) { + final boolean useIsolatedNodePool) { this.jobId = jobId; this.attempt = attempt; this.imageName = imageName; @@ -74,7 +74,6 @@ public AirbyteIntegrationLauncher(final String jobId, this.useIsolatedNodePool = useIsolatedNodePool; } - @Trace(operationName = WORKER_OPERATION_NAME) @Override public Process spec(final Path jobRoot) throws WorkerException { diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index 6bce335ac17a..e0516a162453 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -37,8 +37,6 @@ public class KubeProcessFactory implements ProcessFactory { private final String processRunnerHost; private final boolean isOrchestrator; - private boolean useIsolatedNodePool; - /** * Sets up a process factory with the default processRunnerHost. */ diff --git a/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java b/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java index a67170d82982..5617f8a14175 100644 --- a/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java +++ b/airbyte-commons-worker/src/test/java/io/airbyte/workers/process/AirbyteIntegrationLauncherTest.java @@ -122,7 +122,7 @@ void read() throws WorkerException { void write() throws WorkerException { launcher.write(JOB_ROOT, CONFIG, "{}", CATALOG, "{}"); - Mockito.verify(processFactory).create(WRITE_STEP, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, true, CONFIG_CATALOG_FILES, null, + Mockito.verify(processFactory).create(WRITE_STEP, JOB_ID, JOB_ATTEMPT, JOB_ROOT, FAKE_IMAGE, false, true, CONFIG_CATALOG_FILES, null, workerConfigs.getResourceRequirements(), Map.of(JOB_TYPE, SYNC_JOB, SYNC_STEP, WRITE_STEP), JOB_METADATA, diff --git a/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml index 818cd8d24b18..3ec7819bfca0 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobCheckConnectionConfig.yaml @@ -20,4 +20,4 @@ properties: existingJavaType: io.airbyte.commons.version.Version isCustomConnector: description: determine if the running image is a custom connector. - type: boolean \ No newline at end of file + type: boolean diff --git a/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml index 507dfafd9ee6..af3162a4b1ab 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobGetSpecConfig.yaml @@ -12,4 +12,4 @@ properties: type: string isCustomConnector: description: determine if the running image is a custom connector. - type: boolean \ No newline at end of file + type: boolean diff --git a/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml index 6c2529fad220..3bba82085306 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobResetConnectionConfig.yaml @@ -57,4 +57,4 @@ properties: type: boolean isDestinationCustomConnector: description: determine if the running image of the destination is a custom connector. - type: boolean \ No newline at end of file + type: boolean diff --git a/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml b/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml index 7bb14d29176d..1862f17f884c 100644 --- a/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml +++ b/airbyte-config/config-models/src/main/resources/types/JobSyncConfig.yaml @@ -76,4 +76,4 @@ properties: type: boolean isDestinationCustomConnector: description: determine if the destination running image is a custom connector. - type: boolean \ No newline at end of file + type: boolean diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index db7683c390d7..79d6eeb7c6aa 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -192,7 +192,8 @@ public CheckConnectionRead checkDestinationConnectionFromDestinationId(final Des final String imageName = DockerUtils.getTaggedImageName(destinationDef.getDockerRepository(), destinationDef.getDockerImageTag()); final boolean isCustomConnector = destinationDef.getCustom(); final Version protocolVersion = new Version(destinationDef.getProtocolVersion()); - return reportConnectionStatus(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion,isCustomConnector)); + return reportConnectionStatus( + synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion, isCustomConnector)); } public CheckConnectionRead checkDestinationConnectionFromDestinationCreate(final DestinationCoreConfig destinationConfig) @@ -211,7 +212,8 @@ public CheckConnectionRead checkDestinationConnectionFromDestinationCreate(final .withWorkspaceId(destinationConfig.getWorkspaceId()); final String imageName = DockerUtils.getTaggedImageName(destDef.getDockerRepository(), destDef.getDockerImageTag()); final Version protocolVersion = new Version(destDef.getProtocolVersion()); - return reportConnectionStatus(synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion, isCustomConnector)); + return reportConnectionStatus( + synchronousSchedulerClient.createDestinationCheckConnectionJob(destination, imageName, protocolVersion, isCustomConnector)); } public CheckConnectionRead checkDestinationConnectionFromDestinationIdForUpdate(final DestinationUpdate destinationUpdate) @@ -244,7 +246,8 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceId(final Source final boolean bustActorCatalogCache = discoverSchemaRequestBody.getDisableCache() != null && discoverSchemaRequestBody.getDisableCache(); if (currentCatalog.isEmpty() || bustActorCatalogCache) { final SynchronousResponse persistedCatalogId = - synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, connectorVersion, new Version(sourceDef.getProtocolVersion()), isCustomConnector); + synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, connectorVersion, new Version(sourceDef.getProtocolVersion()), + isCustomConnector); final SourceDiscoverSchemaRead discoveredSchema = retrieveDiscoveredSchema(persistedCatalogId); if (discoverSchemaRequestBody.getConnectionId() != null) { diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java index e4702e7882b5..392b792fb0fe 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java @@ -218,7 +218,8 @@ public SourceDefinitionRead createCustomSourceDefinition(final CustomSourceDefin private StandardSourceDefinition sourceDefinitionFromCreate(final SourceDefinitionCreate sourceDefinitionCreate) throws IOException { - final ConnectorSpecification spec = getSpecForImage(sourceDefinitionCreate.getDockerRepository(), sourceDefinitionCreate.getDockerImageTag(), true); + final ConnectorSpecification spec = + getSpecForImage(sourceDefinitionCreate.getDockerRepository(), sourceDefinitionCreate.getDockerImageTag(), true); final Version airbyteProtocolVersion = AirbyteProtocolVersion.getWithDefault(spec.getProtocolVersion()); @@ -247,7 +248,8 @@ public SourceDefinitionRead updateSourceDefinition(final SourceDefinitionUpdate final boolean specNeedsUpdate = !currentSourceDefinition.getDockerImageTag().equals(sourceDefinitionUpdate.getDockerImageTag()) || sourceDefinitionUpdate.getDockerImageTag().equals(DEV_IMAGE_TAG); final ConnectorSpecification spec = specNeedsUpdate - ? getSpecForImage(currentSourceDefinition.getDockerRepository(), sourceDefinitionUpdate.getDockerImageTag(), currentSourceDefinition.getCustom()) + ? getSpecForImage(currentSourceDefinition.getDockerRepository(), sourceDefinitionUpdate.getDockerImageTag(), + currentSourceDefinition.getCustom()) : currentSourceDefinition.getSpec(); final ActorDefinitionResourceRequirements updatedResourceReqs = sourceDefinitionUpdate.getResourceRequirements() != null ? ApiPojoConverters.actorDefResourceReqsToInternal(sourceDefinitionUpdate.getResourceRequirements()) @@ -317,7 +319,8 @@ public void deleteCustomSourceDefinition(final SourceDefinitionIdWithWorkspaceId deleteSourceDefinition(new SourceDefinitionIdRequestBody().sourceDefinitionId(definitionId)); } - private ConnectorSpecification getSpecForImage(final String dockerRepository, final String imageTag, final boolean isCustomConnector) throws IOException { + private ConnectorSpecification getSpecForImage(final String dockerRepository, final String imageTag, final boolean isCustomConnector) + throws IOException { final String imageName = DockerUtils.getTaggedImageName(dockerRepository, imageTag); final SynchronousResponse getSpecResponse = schedulerSynchronousClient.createGetSpecJob(imageName, isCustomConnector); return SpecFetcher.getSpecFromJob(getSpecResponse); diff --git a/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java b/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java index 0d250cf2c98b..6ab0bacac3f3 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java +++ b/airbyte-server/src/main/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClient.java @@ -63,7 +63,7 @@ public DefaultSynchronousSchedulerClient(final TemporalClient temporalClient, public SynchronousResponse createSourceCheckConnectionJob(final SourceConnection source, final String dockerImage, final Version protocolVersion, - final boolean isCustomConnector) + final boolean isCustomConnector) throws IOException { final JsonNode sourceConfiguration = oAuthConfigSupplier.injectSourceOAuthParameters( source.getSourceDefinitionId(), @@ -91,7 +91,7 @@ public SynchronousResponse createSourceCheckConne public SynchronousResponse createDestinationCheckConnectionJob(final DestinationConnection destination, final String dockerImage, final Version protocolVersion, - final boolean isCustomConnector) + final boolean isCustomConnector) throws IOException { final JsonNode destinationConfiguration = oAuthConfigSupplier.injectDestinationOAuthParameters( destination.getDestinationDefinitionId(), @@ -120,7 +120,7 @@ public SynchronousResponse createDiscoverSchemaJob(final SourceConnection final String dockerImage, final String connectorVersion, final Version protocolVersion, - final boolean isCustomConnector) + final boolean isCustomConnector) throws IOException { final JsonNode sourceConfiguration = oAuthConfigSupplier.injectSourceOAuthParameters( source.getSourceDefinitionId(), diff --git a/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java b/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java index e794fbf54579..7ff4257e1f2b 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java +++ b/airbyte-server/src/main/java/io/airbyte/server/scheduler/SynchronousSchedulerClient.java @@ -20,17 +20,23 @@ public interface SynchronousSchedulerClient { SynchronousResponse createSourceCheckConnectionJob(SourceConnection source, String dockerImage, - Version protocolVersion,boolean isCustomConnector) + Version protocolVersion, + boolean isCustomConnector) throws IOException; SynchronousResponse createDestinationCheckConnectionJob(DestinationConnection destination, String dockerImage, - Version protocolVersion,boolean isCustomConnector) + Version protocolVersion, + boolean isCustomConnector) throws IOException; - SynchronousResponse createDiscoverSchemaJob(SourceConnection source, String dockerImage, String connectorVersion, Version protocolVersion, boolean isCustomConnector) + SynchronousResponse createDiscoverSchemaJob(SourceConnection source, + String dockerImage, + String connectorVersion, + Version protocolVersion, + boolean isCustomConnector) throws IOException; - SynchronousResponse createGetSpecJob(String dockerImage,boolean isCustomConnector) throws IOException; + SynchronousResponse createGetSpecJob(String dockerImage, boolean isCustomConnector) throws IOException; } diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java index bbe304967dc4..b7269ee76498 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java @@ -412,8 +412,9 @@ void testDiscoverSchemaForSourceFromSourceId() throws IOException, JsonValidatio .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); when(configRepository.getActorCatalog(any(), any(), any())).thenReturn(Optional.empty()); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn(discoverResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn(discoverResponse); final SourceDiscoverSchemaRead actual = schedulerHandler.discoverSchemaForSourceFromSourceId(request); @@ -423,7 +424,8 @@ void testDiscoverSchemaForSourceFromSourceId() throws IOException, JsonValidatio assertTrue(actual.getJobInfo().getSucceeded()); verify(configRepository).getSourceConnection(source.getSourceId()); verify(configRepository).getActorCatalog(eq(request.getSourceId()), eq(SOURCE_DOCKER_TAG), any()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false); } @Test @@ -450,8 +452,9 @@ void testDiscoverSchemaForSourceFromSourceIdCachedCatalog() throws IOException, .withCatalogHash("") .withId(thisCatalogId); when(configRepository.getActorCatalog(any(), any(), any())).thenReturn(Optional.of(actorCatalog)); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn(discoverResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn(discoverResponse); final SourceDiscoverSchemaRead actual = schedulerHandler.discoverSchemaForSourceFromSourceId(request); @@ -491,8 +494,9 @@ void testDiscoverSchemaForSourceFromSourceIdDisableCache() throws IOException, J .withCatalogHash("") .withId(discoveredCatalogId); when(configRepository.getActorCatalogById(discoveredCatalogId)).thenReturn(actorCatalog); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn(discoverResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn(discoverResponse); final SourceDiscoverSchemaRead actual = schedulerHandler.discoverSchemaForSourceFromSourceId(request); @@ -501,7 +505,8 @@ void testDiscoverSchemaForSourceFromSourceIdDisableCache() throws IOException, J assertTrue(actual.getJobInfo().getSucceeded()); verify(configRepository).getSourceConnection(source.getSourceId()); verify(configRepository).getActorCatalog(eq(request.getSourceId()), any(), any()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false); } @Test @@ -516,8 +521,9 @@ void testDiscoverSchemaForSourceFromSourceIdFailed() throws IOException, JsonVal .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn((SynchronousResponse) jobResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn((SynchronousResponse) jobResponse); when(completedJob.getSuccessOutput()).thenReturn(Optional.empty()); when(completedJob.getStatus()).thenReturn(JobStatus.FAILED); @@ -527,7 +533,8 @@ void testDiscoverSchemaForSourceFromSourceIdFailed() throws IOException, JsonVal assertNotNull(actual.getJobInfo()); assertFalse(actual.getJobInfo().getSucceeded()); verify(configRepository).getSourceConnection(source.getSourceId()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false); } @Test @@ -548,8 +555,9 @@ void testDiscoverSchemaFromSourceIdWithConnectionIdNonBreaking() throws IOExcept .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn(discoverResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn(discoverResponse); when(discoverResponse.isSuccess()).thenReturn(true); when(discoverResponse.getOutput()).thenReturn(discoveredCatalogId); @@ -596,8 +604,9 @@ void testDiscoverSchemaFromSourceIdWithConnectionIdBreaking() throws IOException .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); when(configRepository.getSourceConnection(source.getSourceId())).thenReturn(source); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn(discoverResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn(discoverResponse); when(discoverResponse.isSuccess()).thenReturn(true); when(discoverResponse.getOutput()).thenReturn(discoveredCatalogId); @@ -655,8 +664,9 @@ void testDiscoverSchemaForSourceFromSourceCreate() throws JsonValidationExceptio .withDockerImageTag(SOURCE_DOCKER_TAG) .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn(discoverResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn(discoverResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(source.getConfiguration()), any())).thenReturn(source.getConfiguration()); @@ -667,7 +677,8 @@ void testDiscoverSchemaForSourceFromSourceCreate() throws JsonValidationExceptio assertNotNull(actual.getJobInfo()); assertEquals(actual.getCatalogId(), discoverResponse.getOutput()); assertTrue(actual.getJobInfo().getSucceeded()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false); } @Test @@ -687,8 +698,9 @@ void testDiscoverSchemaForSourceFromSourceCreateFailed() throws JsonValidationEx .withDockerImageTag(SOURCE_DOCKER_TAG) .withProtocolVersion(SOURCE_PROTOCOL_VERSION) .withSourceDefinitionId(source.getSourceDefinitionId())); - when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false)) - .thenReturn((SynchronousResponse) jobResponse); + when(synchronousSchedulerClient.createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false)) + .thenReturn((SynchronousResponse) jobResponse); when(secretsRepositoryWriter.statefulSplitEphemeralSecrets( eq(source.getConfiguration()), any())).thenReturn(source.getConfiguration()); @@ -700,7 +712,8 @@ void testDiscoverSchemaForSourceFromSourceCreateFailed() throws JsonValidationEx assertNull(actual.getCatalog()); assertNotNull(actual.getJobInfo()); assertFalse(actual.getJobInfo().getSucceeded()); - verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), false); + verify(synchronousSchedulerClient).createDiscoverSchemaJob(source, SOURCE_DOCKER_IMAGE, SOURCE_DOCKER_TAG, new Version(SOURCE_PROTOCOL_VERSION), + false); } @Test diff --git a/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java b/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java index 25586aa1b89c..909c21a4b71b 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/scheduler/DefaultSynchronousSchedulerClientTest.java @@ -204,8 +204,7 @@ void testCreateSourceCheckConnectionJob() throws IOException { final JobCheckConnectionConfig jobCheckConnectionConfig = new JobCheckConnectionConfig() .withConnectionConfiguration(SOURCE_CONNECTION.getConfiguration()) .withDockerImage(DOCKER_IMAGE) - .withProtocolVersion(PROTOCOL_VERSION) .withIsCustomConnector(false); - + .withProtocolVersion(PROTOCOL_VERSION).withIsCustomConnector(false); final StandardCheckConnectionOutput mockOutput = mock(StandardCheckConnectionOutput.class); final ConnectorJobOutput jobOutput = new ConnectorJobOutput().withCheckConnection(mockOutput); diff --git a/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml b/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml index 5dcf6d11af0c..bcdb7fe3a489 100644 --- a/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml +++ b/airbyte-worker-models/src/main/resources/workers_models/IntegrationLauncherConfig.yaml @@ -20,4 +20,4 @@ properties: type: object existingJavaType: io.airbyte.commons.version.Version isCustomConnector: - type: boolean \ No newline at end of file + type: boolean diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index 07827acf7180..bf6c73169646 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -85,6 +85,12 @@ public Map specJobKubeAnnotations(@Value("${airbyte.worker.spec. return splitKVPairsFromEnvString(kubeAnnotations); } + @Singleton + @Named("customNodeSelectors") + public Map customNodeSelectors(@Value("${airbyte.worker.custom.kube.node-selectors}") final String kubeNodeSelectors) { + return splitKVPairsFromEnvString(kubeNodeSelectors); + } + @Singleton @Named("specJobKubeNodeSelectors") public Map specJobKubeNodeSelectors(@Value("${airbyte.worker.spec.kube.node-selectors}") final String kubeNodeSelectors) { @@ -206,7 +212,7 @@ public WorkerConfigs checkWorkerConfigs( @Named("checkResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("checkJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, + @Named("customNodeSelectors") final Map customNodeSelectors, @Named("checkJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -238,9 +244,9 @@ public WorkerConfigs defaultWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, + @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("defaultJobKubeAnnotations") final Map annotations, + @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -272,9 +278,9 @@ public WorkerConfigs discoverWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("discoverJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, + @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("discoverJobKubeAnnotations") final Map annotations, + @Named("discoverJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -305,9 +311,9 @@ public WorkerConfigs replicationWorkerConfigs( @Named("replicationResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, + @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("defaultJobKubeAnnotations") final Map annotations, + @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, @@ -339,9 +345,9 @@ public WorkerConfigs specWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("specJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, + @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("specJobKubeAnnotations") final Map annotations, + @Named("specJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java index b90cf33c2347..280196052ddb 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java @@ -83,7 +83,8 @@ public ConnectorJobOutput run(final JobRunConfig jobRunConfig, final Integration ApmTraceUtils.addTagsToTrace(Map.of(ATTEMPT_NUMBER_KEY, jobRunConfig.getAttemptId(), DOCKER_IMAGE_KEY, launcherConfig.getDockerImage(), JOB_ID_KEY, jobRunConfig.getJobId())); - final Supplier inputSupplier = () -> new JobGetSpecConfig().withDockerImage(launcherConfig.getDockerImage()).withIsCustomConnector(launcherConfig.getIsCustomConnector()); + final Supplier inputSupplier = + () -> new JobGetSpecConfig().withDockerImage(launcherConfig.getDockerImage()).withIsCustomConnector(launcherConfig.getIsCustomConnector()); final ActivityExecutionContext context = Activity.getExecutionContext(); diff --git a/airbyte-workers/src/main/resources/application.yml b/airbyte-workers/src/main/resources/application.yml index 78887a196f87..acf800ef0811 100644 --- a/airbyte-workers/src/main/resources/application.yml +++ b/airbyte-workers/src/main/resources/application.yml @@ -102,6 +102,9 @@ airbyte: memory: limit: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_LIMIT:} request: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_REQUEST:} + custom: + kube: + node-selectors: ${CUSTOM_JOB_KUBE_NODE_SELECTORS:} connection: enabled: ${SHOULD_RUN_CONNECTION_MANAGER_WORKFLOWS:true} discover: From 1b65f74b93fb2f01b71ccfa84712087e80f24f4b Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 22 Nov 2022 15:50:00 -0800 Subject: [PATCH 04/18] fix more missing configs pass along --- .../io/airbyte/commons/temporal/TemporalClient.java | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java b/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java index 8073ff001cb8..3c8ec42e4985 100644 --- a/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java +++ b/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/TemporalClient.java @@ -347,7 +347,8 @@ public TemporalResponse submitCheckConnection(final UUID job .withJobId(jobId.toString()) .withAttemptId((long) attempt) .withDockerImage(config.getDockerImage()) - .withProtocolVersion(config.getProtocolVersion()); + .withProtocolVersion(config.getProtocolVersion()) + .withIsCustomConnector(config.getIsCustomConnector()); final StandardCheckConnectionInput input = new StandardCheckConnectionInput().withConnectionConfiguration(config.getConnectionConfiguration()); return execute(jobRunConfig, @@ -362,7 +363,8 @@ public TemporalResponse submitDiscoverSchema(final UUID jobI .withJobId(jobId.toString()) .withAttemptId((long) attempt) .withDockerImage(config.getDockerImage()) - .withProtocolVersion(config.getProtocolVersion()); + .withProtocolVersion(config.getProtocolVersion()) + .withIsCustomConnector(config.getIsCustomConnector()); final StandardDiscoverCatalogInput input = new StandardDiscoverCatalogInput().withConnectionConfiguration(config.getConnectionConfiguration()) .withSourceId(config.getSourceId()).withConnectorVersion(config.getConnectorVersion()).withConfigHash(config.getConfigHash()); @@ -377,13 +379,15 @@ public TemporalResponse submitSync(final long jobId, final i .withJobId(String.valueOf(jobId)) .withAttemptId((long) attempt) .withDockerImage(config.getSourceDockerImage()) - .withProtocolVersion(config.getSourceProtocolVersion()); + .withProtocolVersion(config.getSourceProtocolVersion()) + .withIsCustomConnector(config.getIsSourceCustomConnector()); final IntegrationLauncherConfig destinationLauncherConfig = new IntegrationLauncherConfig() .withJobId(String.valueOf(jobId)) .withAttemptId((long) attempt) .withDockerImage(config.getDestinationDockerImage()) - .withProtocolVersion(config.getDestinationProtocolVersion()); + .withProtocolVersion(config.getDestinationProtocolVersion()) + .withIsCustomConnector(config.getIsDestinationCustomConnector()); final StandardSyncInput input = new StandardSyncInput() .withNamespaceDefinition(config.getNamespaceDefinition()) From cf5e2b3055b354d7112147fff82e1d654f1f9d83 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 22 Nov 2022 17:19:01 -0800 Subject: [PATCH 05/18] fix sync config --- .../persistence/job/DefaultJobCreator.java | 15 +++++++++------ .../io/airbyte/persistence/job/JobCreator.java | 7 ++++--- .../job/factory/DefaultSyncJobFactory.java | 11 +++-------- .../job/factory/DefaultSyncJobFactoryTest.java | 17 +++++++++++------ 4 files changed, 27 insertions(+), 23 deletions(-) diff --git a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/DefaultJobCreator.java b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/DefaultJobCreator.java index 82d108dcce52..553b5b8dadab 100644 --- a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/DefaultJobCreator.java +++ b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/DefaultJobCreator.java @@ -6,7 +6,6 @@ import com.fasterxml.jackson.databind.JsonNode; import io.airbyte.commons.version.Version; -import io.airbyte.config.ActorDefinitionResourceRequirements; import io.airbyte.config.DestinationConnection; import io.airbyte.config.JobConfig; import io.airbyte.config.JobConfig.ConfigType; @@ -16,6 +15,8 @@ import io.airbyte.config.ResetSourceConfiguration; import io.airbyte.config.ResourceRequirements; import io.airbyte.config.SourceConnection; +import io.airbyte.config.StandardDestinationDefinition; +import io.airbyte.config.StandardSourceDefinition; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSyncOperation; import io.airbyte.config.State; @@ -58,8 +59,8 @@ public Optional createSyncJob(final SourceConnection source, final Version destinationProtocolVersion, final List standardSyncOperations, @Nullable final JsonNode webhookOperationConfigs, - @Nullable final ActorDefinitionResourceRequirements sourceResourceReqs, - @Nullable final ActorDefinitionResourceRequirements destinationResourceReqs) + final StandardSourceDefinition sourceDefinition, + final StandardDestinationDefinition destinationDefinition) throws IOException { // reusing this isn't going to quite work. @@ -68,12 +69,12 @@ public Optional createSyncJob(final SourceConnection source, workerResourceRequirements); final ResourceRequirements mergedSrcResourceReq = ResourceRequirementsUtils.getResourceRequirements( standardSync.getResourceRequirements(), - sourceResourceReqs, + sourceDefinition.getResourceRequirements(), workerResourceRequirements, JobType.SYNC); final ResourceRequirements mergedDstResourceReq = ResourceRequirementsUtils.getResourceRequirements( standardSync.getResourceRequirements(), - destinationResourceReqs, + destinationDefinition.getResourceRequirements(), workerResourceRequirements, JobType.SYNC); @@ -93,7 +94,9 @@ public Optional createSyncJob(final SourceConnection source, .withState(null) .withResourceRequirements(mergedOrchestratorResourceReq) .withSourceResourceRequirements(mergedSrcResourceReq) - .withDestinationResourceRequirements(mergedDstResourceReq); + .withDestinationResourceRequirements(mergedDstResourceReq) + .withIsSourceCustomConnector(sourceDefinition.getCustom()) + .withIsDestinationCustomConnector(destinationDefinition.getCustom()); getCurrentConnectionState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); diff --git a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/JobCreator.java b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/JobCreator.java index 687d795dc3f4..df900abd6287 100644 --- a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/JobCreator.java +++ b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/JobCreator.java @@ -6,9 +6,10 @@ import com.fasterxml.jackson.databind.JsonNode; import io.airbyte.commons.version.Version; -import io.airbyte.config.ActorDefinitionResourceRequirements; import io.airbyte.config.DestinationConnection; import io.airbyte.config.SourceConnection; +import io.airbyte.config.StandardDestinationDefinition; +import io.airbyte.config.StandardSourceDefinition; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSyncOperation; import io.airbyte.protocol.models.StreamDescriptor; @@ -38,8 +39,8 @@ Optional createSyncJob(SourceConnection source, Version destinationProtocolVersion, List standardSyncOperations, @Nullable JsonNode webhookOperationConfigs, - @Nullable ActorDefinitionResourceRequirements sourceResourceReqs, - @Nullable ActorDefinitionResourceRequirements destinationResourceReqs) + StandardSourceDefinition sourceDefinition, + StandardDestinationDefinition destinationDefinition) throws IOException; /** diff --git a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java index b83d8f3497dc..8ee1d21b05b7 100644 --- a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java +++ b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java @@ -8,7 +8,6 @@ import com.google.common.collect.Lists; import io.airbyte.commons.docker.DockerUtils; import io.airbyte.commons.version.Version; -import io.airbyte.config.ActorDefinitionResourceRequirements; import io.airbyte.config.DestinationConnection; import io.airbyte.config.SourceConnection; import io.airbyte.config.StandardDestinationDefinition; @@ -78,13 +77,10 @@ public Long create(final UUID connectionId) { standardSyncOperations.add(standardSyncOperation); } - ActorDefinitionResourceRequirements sourceResourceRequirements = sourceDefinition.getResourceRequirements(); - ActorDefinitionResourceRequirements destinationResourceRequirements = destinationDefinition.getResourceRequirements(); - // for OSS users, make it possible to ignore default actor-level resource requirements if (!connectorSpecificResourceDefaultsEnabled) { - sourceResourceRequirements = null; - destinationResourceRequirements = null; + sourceDefinition.setResourceRequirements(null); + destinationDefinition.setResourceRequirements(null); } return jobCreator.createSyncJob( @@ -97,8 +93,7 @@ public Long create(final UUID connectionId) { new Version(destinationDefinition.getProtocolVersion()), standardSyncOperations, workspace.getWebhookOperationConfigs(), - sourceResourceRequirements, - destinationResourceRequirements) + sourceDefinition, destinationDefinition) .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-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactoryTest.java b/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactoryTest.java index dc2cfdc9285d..120f9e4b984f 100644 --- a/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactoryTest.java +++ b/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactoryTest.java @@ -63,6 +63,7 @@ void createSyncJobFromConnectionId() throws JsonValidationException, ConfigNotFo final SourceConnection sourceConnection = new SourceConnection().withSourceDefinitionId(sourceDefinitionId); final DestinationConnection destinationConnection = new DestinationConnection().withDestinationDefinitionId(destinationDefinitionId); + final String srcDockerRepo = "srcrepo"; final String srcDockerTag = "tag"; final String srcDockerImage = DockerUtils.getTaggedImageName(srcDockerRepo, srcDockerTag); @@ -72,6 +73,12 @@ void createSyncJobFromConnectionId() throws JsonValidationException, ConfigNotFo final String dstDockerTag = "tag"; final String dstDockerImage = DockerUtils.getTaggedImageName(dstDockerRepo, dstDockerTag); final Version dstProtocolVersion = new Version("0.3.2"); + final StandardSourceDefinition standardSourceDefinition = + new StandardSourceDefinition().withSourceDefinitionId(sourceDefinitionId).withDockerRepository(srcDockerRepo) + .withDockerImageTag(srcDockerTag).withProtocolVersion(srcProtocolVersion.serialize()); + final StandardDestinationDefinition standardDestinationDefinition = + new StandardDestinationDefinition().withDestinationDefinitionId(destinationDefinitionId).withDockerRepository(dstDockerRepo) + .withDockerImageTag(dstDockerTag).withProtocolVersion(dstProtocolVersion.serialize()); when(configRepository.getStandardSync(connectionId)).thenReturn(standardSync); when(configRepository.getSourceConnection(sourceId)).thenReturn(sourceConnection); @@ -80,15 +87,13 @@ void createSyncJobFromConnectionId() throws JsonValidationException, ConfigNotFo when( jobCreator.createSyncJob(sourceConnection, destinationConnection, standardSync, srcDockerImage, srcProtocolVersion, dstDockerImage, dstProtocolVersion, operations, - persistedWebhookConfigs, null, null)) + persistedWebhookConfigs, standardSourceDefinition, standardDestinationDefinition)) .thenReturn(Optional.of(jobId)); when(configRepository.getStandardSourceDefinition(sourceDefinitionId)) - .thenReturn(new StandardSourceDefinition().withSourceDefinitionId(sourceDefinitionId).withDockerRepository(srcDockerRepo) - .withDockerImageTag(srcDockerTag).withProtocolVersion(srcProtocolVersion.serialize())); + .thenReturn(standardSourceDefinition); when(configRepository.getStandardDestinationDefinition(destinationDefinitionId)) - .thenReturn(new StandardDestinationDefinition().withDestinationDefinitionId(destinationDefinitionId).withDockerRepository(dstDockerRepo) - .withDockerImageTag(dstDockerTag).withProtocolVersion(dstProtocolVersion.serialize())); + .thenReturn(standardDestinationDefinition); when(configRepository.getStandardWorkspaceNoSecrets(any(), eq(true))).thenReturn( new StandardWorkspace().withWebhookOperationConfigs(persistedWebhookConfigs)); @@ -100,7 +105,7 @@ void createSyncJobFromConnectionId() throws JsonValidationException, ConfigNotFo verify(jobCreator) .createSyncJob(sourceConnection, destinationConnection, standardSync, srcDockerImage, srcProtocolVersion, dstDockerImage, dstProtocolVersion, operations, persistedWebhookConfigs, - null, null); + standardSourceDefinition, standardDestinationDefinition); } } From ccef1d81f9534e782866b9c94427240388cfc99c Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 22 Nov 2022 21:46:01 -0800 Subject: [PATCH 06/18] test fix --- .../job/DefaultJobCreatorTest.java | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/DefaultJobCreatorTest.java b/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/DefaultJobCreatorTest.java index 8fedbae6cd4b..18acdba973b0 100644 --- a/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/DefaultJobCreatorTest.java +++ b/airbyte-persistence/job-persistence/src/test/java/io/airbyte/persistence/job/DefaultJobCreatorTest.java @@ -29,6 +29,8 @@ import io.airbyte.config.ResetSourceConfiguration; import io.airbyte.config.ResourceRequirements; import io.airbyte.config.SourceConnection; +import io.airbyte.config.StandardDestinationDefinition; +import io.airbyte.config.StandardSourceDefinition; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSyncOperation; import io.airbyte.config.StandardSyncOperation.OperatorType; @@ -69,6 +71,9 @@ class DefaultJobCreatorTest { private static final DestinationConnection DESTINATION_CONNECTION; private static final StandardSync STANDARD_SYNC; private static final StandardSyncOperation STANDARD_SYNC_OPERATION; + + private static final StandardSourceDefinition STANDARD_SOURCE_DEFINITION; + private static final StandardDestinationDefinition STANDARD_DESTINATION_DEFINITION; private static final long JOB_ID = 12L; private JobPersistence jobPersistence; @@ -149,6 +154,9 @@ class DefaultJobCreatorTest { PERSISTED_WEBHOOK_CONFIGS = Jsons.deserialize( String.format("{\"webhookConfigs\": [{\"id\": \"%s\", \"name\": \"%s\", \"authToken\": {\"_secret\": \"a-secret_v1\"}}]}", WEBHOOK_CONFIG_ID, WEBHOOK_NAME)); + + STANDARD_SOURCE_DEFINITION = new StandardSourceDefinition().withCustom(false); + STANDARD_DESTINATION_DEFINITION = new StandardDestinationDefinition().withCustom(false); } @BeforeEach @@ -180,7 +188,9 @@ void testCreateSyncJob() throws IOException { .withResourceRequirements(workerResourceRequirements) .withSourceResourceRequirements(workerResourceRequirements) .withDestinationResourceRequirements(workerResourceRequirements) - .withWebhookOperationConfigs(PERSISTED_WEBHOOK_CONFIGS); + .withWebhookOperationConfigs(PERSISTED_WEBHOOK_CONFIGS) + .withIsSourceCustomConnector(false) + .withIsDestinationCustomConnector(false); final JobConfig jobConfig = new JobConfig() .withConfigType(JobConfig.ConfigType.SYNC) @@ -199,8 +209,8 @@ void testCreateSyncJob() throws IOException { DESTINATION_PROTOCOL_VERSION, List.of(STANDARD_SYNC_OPERATION), PERSISTED_WEBHOOK_CONFIGS, - null, - null).orElseThrow(); + STANDARD_SOURCE_DEFINITION, + STANDARD_DESTINATION_DEFINITION).orElseThrow(); assertEquals(JOB_ID, jobId); } @@ -237,8 +247,7 @@ void testCreateSyncJobEnsureNoQueuing() throws IOException { DESTINATION_PROTOCOL_VERSION, List.of(STANDARD_SYNC_OPERATION), null, - null, - null).isEmpty()); + STANDARD_SOURCE_DEFINITION, STANDARD_DESTINATION_DEFINITION).isEmpty()); } @Test @@ -253,8 +262,7 @@ void testCreateSyncJobDefaultWorkerResourceReqs() throws IOException { DESTINATION_PROTOCOL_VERSION, List.of(STANDARD_SYNC_OPERATION), null, - null, - null); + STANDARD_SOURCE_DEFINITION, STANDARD_DESTINATION_DEFINITION); final JobSyncConfig expectedJobSyncConfig = new JobSyncConfig() .withNamespaceDefinition(STANDARD_SYNC.getNamespaceDefinition()) @@ -270,7 +278,9 @@ void testCreateSyncJobDefaultWorkerResourceReqs() throws IOException { .withOperationSequence(List.of(STANDARD_SYNC_OPERATION)) .withResourceRequirements(workerResourceRequirements) .withSourceResourceRequirements(workerResourceRequirements) - .withDestinationResourceRequirements(workerResourceRequirements); + .withDestinationResourceRequirements(workerResourceRequirements) + .withIsSourceCustomConnector(false) + .withIsDestinationCustomConnector(false); final JobConfig expectedJobConfig = new JobConfig() .withConfigType(JobConfig.ConfigType.SYNC) @@ -300,8 +310,7 @@ void testCreateSyncJobConnectionResourceReqs() throws IOException { DESTINATION_PROTOCOL_VERSION, List.of(STANDARD_SYNC_OPERATION), null, - null, - null); + STANDARD_SOURCE_DEFINITION, STANDARD_DESTINATION_DEFINITION); final JobSyncConfig expectedJobSyncConfig = new JobSyncConfig() .withNamespaceDefinition(STANDARD_SYNC.getNamespaceDefinition()) @@ -317,7 +326,9 @@ void testCreateSyncJobConnectionResourceReqs() throws IOException { .withOperationSequence(List.of(STANDARD_SYNC_OPERATION)) .withResourceRequirements(standardSyncResourceRequirements) .withSourceResourceRequirements(standardSyncResourceRequirements) - .withDestinationResourceRequirements(standardSyncResourceRequirements); + .withDestinationResourceRequirements(standardSyncResourceRequirements) + .withIsSourceCustomConnector(false) + .withIsDestinationCustomConnector(false); final JobConfig expectedJobConfig = new JobConfig() .withConfigType(JobConfig.ConfigType.SYNC) @@ -351,9 +362,9 @@ void testCreateSyncJobSourceAndDestinationResourceReqs() throws IOException { DESTINATION_PROTOCOL_VERSION, List.of(STANDARD_SYNC_OPERATION), null, - new ActorDefinitionResourceRequirements().withDefault(sourceResourceRequirements), - new ActorDefinitionResourceRequirements().withJobSpecific(List.of( - new JobTypeResourceLimit().withJobType(JobType.SYNC).withResourceRequirements(destResourceRequirements)))); + new StandardSourceDefinition().withResourceRequirements(new ActorDefinitionResourceRequirements().withDefault(sourceResourceRequirements)), + new StandardDestinationDefinition().withResourceRequirements(new ActorDefinitionResourceRequirements().withJobSpecific(List.of( + new JobTypeResourceLimit().withJobType(JobType.SYNC).withResourceRequirements(destResourceRequirements))))); final JobSyncConfig expectedJobSyncConfig = new JobSyncConfig() .withNamespaceDefinition(STANDARD_SYNC.getNamespaceDefinition()) @@ -369,7 +380,9 @@ void testCreateSyncJobSourceAndDestinationResourceReqs() throws IOException { .withOperationSequence(List.of(STANDARD_SYNC_OPERATION)) .withResourceRequirements(workerResourceRequirements) .withSourceResourceRequirements(sourceResourceRequirements) - .withDestinationResourceRequirements(destResourceRequirements); + .withDestinationResourceRequirements(destResourceRequirements) + .withIsSourceCustomConnector(false) + .withIsDestinationCustomConnector(false); final JobConfig expectedJobConfig = new JobConfig() .withConfigType(JobConfig.ConfigType.SYNC) From 6adaddcc34bea50380664cda2d5f5f92ca0736be Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Wed, 23 Nov 2022 11:39:11 -0800 Subject: [PATCH 07/18] add a flag to specify if we want to use a separate pool or not --- .../main/java/io/airbyte/workers/WorkerConfigs.java | 10 +++++----- .../src/main/java/io/airbyte/config/Configs.java | 5 +++++ .../src/main/java/io/airbyte/config/EnvConfigs.java | 6 ++++++ airbyte-workers/src/main/resources/application.yml | 3 --- 4 files changed, 16 insertions(+), 8 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java index de3355920a02..e6f064dcf83a 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java @@ -42,7 +42,7 @@ public WorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), configs.getJobKubeNodeSelectors(), - configs.getIsolatedJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : configs.getJobKubeNodeSelectors(), configs.getJobKubeAnnotations(), configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -74,7 +74,7 @@ public static WorkerConfigs buildSpecWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, - configs.getIsolatedJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : nodeSelectors, annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -106,7 +106,7 @@ public static WorkerConfigs buildCheckWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getCheckJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, - configs.getIsolatedJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : nodeSelectors, annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -138,7 +138,7 @@ public static WorkerConfigs buildDiscoverWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, - configs.getIsolatedJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : nodeSelectors, annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -159,7 +159,7 @@ public static WorkerConfigs buildReplicationWorkerConfigs(final Configs configs) .withMemoryLimit(configs.getReplicationOrchestratorMemoryLimit()), configs.getJobKubeTolerations(), configs.getJobKubeNodeSelectors(), - configs.getIsolatedJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : configs.getJobKubeNodeSelectors(), configs.getJobKubeAnnotations(), configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java b/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java index 51b8f8c4a8c4..889f0102a598 100644 --- a/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java +++ b/airbyte-config/config-models/src/main/java/io/airbyte/config/Configs.java @@ -410,6 +410,11 @@ public interface Configs { */ Map getIsolatedJobKubeNodeSelectors(); + /** + * Define if we want to run custom connector related jobs in a separate node pool. + */ + boolean getUseCustomKubeNodeSelector(); + /** * Define node selectors for Spec job pods specifically. Each kv-pair is separated by a `,`. */ diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java b/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java index 9928ec62b987..6e464098b109 100644 --- a/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java +++ b/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java @@ -73,6 +73,7 @@ public class EnvConfigs implements Configs { public static final String JOB_KUBE_TOLERATIONS = "JOB_KUBE_TOLERATIONS"; public static final String JOB_KUBE_NODE_SELECTORS = "JOB_KUBE_NODE_SELECTORS"; public static final String JOB_ISOLATED_KUBE_NODE_SELECTORS = "JOB_ISOLATED_KUBE_NODE_SELECTORS"; + public static final String USE_CUSTOM_NODE_SELECTOR = "USE_CUSTOM_NODE_SELECTOR"; public static final String JOB_KUBE_ANNOTATIONS = "JOB_KUBE_ANNOTATIONS"; public static final String JOB_KUBE_SOCAT_IMAGE = "JOB_KUBE_SOCAT_IMAGE"; public static final String JOB_KUBE_BUSYBOX_IMAGE = "JOB_KUBE_BUSYBOX_IMAGE"; @@ -612,6 +613,11 @@ public Map getIsolatedJobKubeNodeSelectors() { return splitKVPairsFromEnvString(getEnvOrDefault(JOB_ISOLATED_KUBE_NODE_SELECTORS, "")); } + @Override + public boolean getUseCustomKubeNodeSelector() { + return getEnvOrDefault(USE_CUSTOM_NODE_SELECTOR, false); + } + /** * Returns a map of node selectors for Spec job pods specifically. * diff --git a/airbyte-workers/src/main/resources/application.yml b/airbyte-workers/src/main/resources/application.yml index acf800ef0811..78887a196f87 100644 --- a/airbyte-workers/src/main/resources/application.yml +++ b/airbyte-workers/src/main/resources/application.yml @@ -102,9 +102,6 @@ airbyte: memory: limit: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_LIMIT:} request: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_REQUEST:} - custom: - kube: - node-selectors: ${CUSTOM_JOB_KUBE_NODE_SELECTORS:} connection: enabled: ${SHOULD_RUN_CONNECTION_MANAGER_WORKFLOWS:true} discover: From 3740aeb6b07f309d514121d746684cc52d7d5490 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Wed, 23 Nov 2022 13:45:55 -0800 Subject: [PATCH 08/18] add missing micronaut configs --- .../WorkerConfigurationBeanFactory.java | 25 ++++++++++++------- .../src/main/resources/application.yml | 4 +++ 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index bf6c73169646..e5b1743830bc 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -91,6 +91,12 @@ public Map customNodeSelectors(@Value("${airbyte.worker.custom.k return splitKVPairsFromEnvString(kubeNodeSelectors); } + @Singleton + @Named("useCustomNodeSelector") + public boolean useCustomNodeSelector(@Value("${airbyte.worker.custom.kube.node-selectors}") final boolean kubeNodeSelectors) { + return kubeNodeSelectors; + } + @Singleton @Named("specJobKubeNodeSelectors") public Map specJobKubeNodeSelectors(@Value("${airbyte.worker.spec.kube.node-selectors}") final String kubeNodeSelectors) { @@ -213,6 +219,7 @@ public WorkerConfigs checkWorkerConfigs( final List jobKubeTolerations, @Named("checkJobKubeNodeSelectors") final Map nodeSelectors, @Named("customNodeSelectors") final Map customNodeSelectors, + @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, @Named("checkJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -226,7 +233,7 @@ public WorkerConfigs checkWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - customNodeSelectors, + useCustomNodeSelector ? customNodeSelectors : nodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -245,7 +252,7 @@ public WorkerConfigs defaultWorkerConfigs( final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, @Named("customNodeSelectors") final Map customNodeSelectors, - + @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -259,7 +266,7 @@ public WorkerConfigs defaultWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - customNodeSelectors, + useCustomNodeSelector ? customNodeSelectors : nodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -279,7 +286,7 @@ public WorkerConfigs discoverWorkerConfigs( final List jobKubeTolerations, @Named("discoverJobKubeNodeSelectors") final Map nodeSelectors, @Named("customNodeSelectors") final Map customNodeSelectors, - + @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, @Named("discoverJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -293,7 +300,7 @@ public WorkerConfigs discoverWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - customNodeSelectors, + useCustomNodeSelector ? customNodeSelectors : nodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -312,7 +319,7 @@ public WorkerConfigs replicationWorkerConfigs( final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, @Named("customNodeSelectors") final Map customNodeSelectors, - + @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -326,7 +333,7 @@ public WorkerConfigs replicationWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - customNodeSelectors, + useCustomNodeSelector ? customNodeSelectors : nodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -346,7 +353,7 @@ public WorkerConfigs specWorkerConfigs( final List jobKubeTolerations, @Named("specJobKubeNodeSelectors") final Map nodeSelectors, @Named("customNodeSelectors") final Map customNodeSelectors, - + @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, @Named("specJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -360,7 +367,7 @@ public WorkerConfigs specWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - customNodeSelectors, + useCustomNodeSelector ? customNodeSelectors : nodeSelectors, annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, diff --git a/airbyte-workers/src/main/resources/application.yml b/airbyte-workers/src/main/resources/application.yml index 78887a196f87..9ec9cd428468 100644 --- a/airbyte-workers/src/main/resources/application.yml +++ b/airbyte-workers/src/main/resources/application.yml @@ -102,6 +102,10 @@ airbyte: memory: limit: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_LIMIT:} request: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_REQUEST:} + custom: + kube: + use-custom-node-selector: ${USE_CUSTOM_NODE_SELECTOR:false} + node-selectors: ${CUSTOM_JOB_KUBE_NODE_SELECTORS:} connection: enabled: ${SHOULD_RUN_CONNECTION_MANAGER_WORKFLOWS:true} discover: From 67a489086509f7070ba391a805d9f64fb9b46a81 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Wed, 23 Nov 2022 16:03:51 -0800 Subject: [PATCH 09/18] make orchestrator job to run in custom pool too --- .../java/io/airbyte/workers/sync/DbtLauncherWorker.java | 3 ++- .../main/java/io/airbyte/workers/sync/LauncherWorker.java | 7 +++++-- .../airbyte/workers/sync/NormalizationLauncherWorker.java | 3 ++- .../io/airbyte/workers/sync/ReplicationLauncherWorker.java | 3 ++- 4 files changed, 11 insertions(+), 5 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java index efad287d22b5..ff1836234b89 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java @@ -43,7 +43,8 @@ public DbtLauncherWorker(final UUID connectionId, activityContext, serverPort, temporalUtils, - workerConfigs); + workerConfigs, + false); } } diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java index 839ca31d4df5..f078ab41de46 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java @@ -73,6 +73,7 @@ public class LauncherWorker implements Worker { private final TemporalUtils temporalUtils; private final WorkerConfigs workerConfigs; + private final boolean isCustomConnector; private final AtomicBoolean cancelled = new AtomicBoolean(false); private AsyncOrchestratorPodProcess process; @@ -87,7 +88,8 @@ public LauncherWorker(final UUID connectionId, final Supplier activityContext, final Integer serverPort, final TemporalUtils temporalUtils, - final WorkerConfigs workerConfigs) { + final WorkerConfigs workerConfigs, + final boolean isCustomConnector) { this.connectionId = connectionId; this.application = application; @@ -101,6 +103,7 @@ public LauncherWorker(final UUID connectionId, this.serverPort = serverPort; this.temporalUtils = temporalUtils; this.workerConfigs = workerConfigs; + this.isCustomConnector = isCustomConnector; } @Trace(operationName = WORKER_OPERATION_NAME) @@ -179,7 +182,7 @@ public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException resourceRequirements, fileMap, portMap, - workerConfigs.getworkerKubeNodeSelectors()); + isCustomConnector ? workerConfigs.getWorkerIsolatedKubeNodeSelectors() : workerConfigs.getworkerKubeNodeSelectors()); } catch (final KubernetesClientException e) { ApmTraceUtils.addExceptionToTrace(e); throw new WorkerException( diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java index 1c47a5435baf..dd274e76ef86 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java @@ -44,7 +44,8 @@ public NormalizationLauncherWorker(final UUID connectionId, activityContext, serverPort, temporalUtils, - workerConfigs); + workerConfigs, + false); } diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/ReplicationLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/ReplicationLauncherWorker.java index da9767b68164..905aed0c6387 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/ReplicationLauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/ReplicationLauncherWorker.java @@ -54,7 +54,8 @@ public ReplicationLauncherWorker(final UUID connectionId, activityContext, serverPort, temporalUtils, - workerConfigs); + workerConfigs, + sourceLauncherConfig.getIsCustomConnector() || destinationLauncherConfig.getIsCustomConnector()); } } From ad955669ff6ca51eb338e70d4ebc86d1e068026e Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Wed, 23 Nov 2022 16:29:12 -0800 Subject: [PATCH 10/18] micronaut fix --- .../airbyte/workers/config/WorkerConfigurationBeanFactory.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index e5b1743830bc..d7daa5bad8c2 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -93,7 +93,7 @@ public Map customNodeSelectors(@Value("${airbyte.worker.custom.k @Singleton @Named("useCustomNodeSelector") - public boolean useCustomNodeSelector(@Value("${airbyte.worker.custom.kube.node-selectors}") final boolean kubeNodeSelectors) { + public boolean useCustomNodeSelector(@Value("${airbyte.worker.custom.kube.use-custom-node-selector}") final boolean kubeNodeSelectors) { return kubeNodeSelectors; } From 5d238a738ed3fa134b7db9ff30fc642d7c08742a Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 29 Nov 2022 08:41:18 -0800 Subject: [PATCH 11/18] pass config to orchestrator --- .../io/airbyte/commons/temporal/sync/OrchestratorConstants.java | 2 ++ .../container_orchestrator/ReplicationJobOrchestrator.java | 1 + 2 files changed, 3 insertions(+) diff --git a/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/sync/OrchestratorConstants.java b/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/sync/OrchestratorConstants.java index 7a7c7806d7d0..877b1a5ddcce 100644 --- a/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/sync/OrchestratorConstants.java +++ b/airbyte-commons-temporal/src/main/java/io/airbyte/commons/temporal/sync/OrchestratorConstants.java @@ -31,6 +31,8 @@ public class OrchestratorConstants { EnvConfigs.JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_SECRET, EnvConfigs.JOB_KUBE_SIDECAR_CONTAINER_IMAGE_PULL_POLICY, EnvConfigs.JOB_KUBE_NODE_SELECTORS, + EnvConfigs.JOB_ISOLATED_KUBE_NODE_SELECTORS, + EnvConfigs.USE_CUSTOM_NODE_SELECTOR, EnvConfigs.DOCKER_NETWORK, EnvConfigs.LOCAL_DOCKER_MOUNT, EnvConfigs.WORKSPACE_DOCKER_MOUNT, diff --git a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java index 6c49bc3c9183..97bf4af3bb6c 100644 --- a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java +++ b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java @@ -96,6 +96,7 @@ public Optional runJob() throws Exception { final IntegrationLauncherConfig destinationLauncherConfig = JobOrchestrator.readAndDeserializeFile( Path.of(KubePodProcess.CONFIG_DIR, ReplicationLauncherWorker.INIT_FILE_DESTINATION_LAUNCHER_CONFIG), IntegrationLauncherConfig.class); + log.info("sourceLauncherConfig is: " + sourceLauncherConfig.toString()); ApmTraceUtils.addTagsToTrace(Map.of(JOB_ID_KEY, jobRunConfig.getJobId(), DESTINATION_DOCKER_IMAGE_KEY, destinationLauncherConfig.getDockerImage(), SOURCE_DOCKER_IMAGE_KEY, sourceLauncherConfig.getDockerImage())); From 0ff533216035875e482643114d4eed194dc919ec Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 29 Nov 2022 10:15:46 -0800 Subject: [PATCH 12/18] fix test --- .../source/AbstractSourceConnectorTest.java | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java b/airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java index 26d786fa784e..d5124a109d57 100644 --- a/airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java +++ b/airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java @@ -144,26 +144,26 @@ public void tearDownInternal() throws Exception { protected ConnectorSpecification runSpec() throws WorkerException { return new DefaultGetSpecWorker( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements())) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements(), false)) .run(new JobGetSpecConfig().withDockerImage(getImageName()), jobRoot).getSpec(); } protected StandardCheckConnectionOutput runCheck() throws Exception { return new DefaultCheckConnectionWorker( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements())) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements(), false)) .run(new StandardCheckConnectionInput().withConnectionConfiguration(getConfig()), jobRoot).getCheckConnection(); } protected String runCheckAndGetStatusAsString(final JsonNode config) throws Exception { return new DefaultCheckConnectionWorker( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements())) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements(), false)) .run(new StandardCheckConnectionInput().withConnectionConfiguration(config), jobRoot).getCheckConnection().getStatus().toString(); } protected UUID runDiscover() throws Exception { final UUID toReturn = new DefaultDiscoverCatalogWorker( mConfigRepository, - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements())) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements(), false)) .run(new StandardDiscoverCatalogInput().withSourceId(SOURCE_ID.toString()).withConnectionConfiguration(getConfig()), jobRoot) .getDiscoverCatalogId(); verify(mConfigRepository).writeActorCatalogFetchEvent(lastPersistedCatalog.capture(), any(), any(), any()); @@ -194,7 +194,7 @@ protected List runRead(final ConfiguredAirbyteCatalog catalog, f .withCatalog(catalog); final AirbyteSource source = new DefaultAirbyteSource( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements())); + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, workerConfigs.getResourceRequirements(), false)); final List messages = new ArrayList<>(); source.start(sourceConfig, jobRoot); while (!source.isFinished()) { @@ -244,8 +244,8 @@ protected ResourceRequirements prepareResourceRequirements(final Map Date: Tue, 29 Nov 2022 12:58:43 -0800 Subject: [PATCH 13/18] destination test fix --- .../destination/DestinationAcceptanceTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airbyte-integrations/bases/standard-destination-test/src/main/java/io/airbyte/integrations/standardtest/destination/DestinationAcceptanceTest.java b/airbyte-integrations/bases/standard-destination-test/src/main/java/io/airbyte/integrations/standardtest/destination/DestinationAcceptanceTest.java index a5e422a49030..68108c18856d 100644 --- a/airbyte-integrations/bases/standard-destination-test/src/main/java/io/airbyte/integrations/standardtest/destination/DestinationAcceptanceTest.java +++ b/airbyte-integrations/bases/standard-destination-test/src/main/java/io/airbyte/integrations/standardtest/destination/DestinationAcceptanceTest.java @@ -1169,13 +1169,13 @@ protected void assertNamespaceNormalization(final String testCaseId, private ConnectorSpecification runSpec() throws WorkerException { return new DefaultGetSpecWorker( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null)) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null, false)) .run(new JobGetSpecConfig().withDockerImage(getImageName()), jobRoot).getSpec(); } protected StandardCheckConnectionOutput runCheck(final JsonNode config) throws WorkerException { return new DefaultCheckConnectionWorker( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null)) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null, false)) .run(new StandardCheckConnectionInput().withConnectionConfiguration(config), jobRoot) .getCheckConnection(); } @@ -1184,7 +1184,7 @@ protected StandardCheckConnectionOutput.Status runCheckWithCatchedException( final JsonNode config) { try { final StandardCheckConnectionOutput standardCheckConnectionOutput = new DefaultCheckConnectionWorker( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null)) + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null, false)) .run(new StandardCheckConnectionInput().withConnectionConfiguration(config), jobRoot) .getCheckConnection(); return standardCheckConnectionOutput.getStatus(); @@ -1196,7 +1196,7 @@ protected StandardCheckConnectionOutput.Status runCheckWithCatchedException( protected AirbyteDestination getDestination() { return new DefaultAirbyteDestination( - new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null)); + new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, null, false)); } protected void runSyncAndVerifyStateOutput(final JsonNode config, From 8616bba49a668658d7c739e83988d956b9009653 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Thu, 1 Dec 2022 16:04:23 -0800 Subject: [PATCH 14/18] PR comments fix --- .../io/airbyte/workers/WorkerConfigs.java | 15 ++++---- .../process/AirbyteIntegrationLauncher.java | 21 ++++++---- .../workers/process/KubeProcessFactory.java | 3 +- .../workers/sync/DbtLauncherWorker.java | 2 + .../airbyte/workers/sync/LauncherWorker.java | 6 ++- .../sync/NormalizationLauncherWorker.java | 2 + .../ReplicationJobOrchestrator.java | 6 ++- .../job/factory/DefaultSyncJobFactory.java | 3 +- .../DestinationDefinitionsHandler.java | 4 +- .../server/handlers/SchedulerHandler.java | 15 ++++++-- .../handlers/SourceDefinitionsHandler.java | 6 ++- .../WorkerConfigurationBeanFactory.java | 38 +++++++++---------- .../src/main/resources/application.yml | 6 +-- 13 files changed, 80 insertions(+), 47 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java index e6f064dcf83a..5d7efb720c2c 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/WorkerConfigs.java @@ -9,6 +9,7 @@ import io.airbyte.config.TolerationPOJO; import java.util.List; import java.util.Map; +import java.util.Optional; import lombok.AllArgsConstructor; @AllArgsConstructor @@ -18,7 +19,7 @@ public class WorkerConfigs { private final ResourceRequirements resourceRequirements; private final List workerKubeTolerations; private final Map workerKubeNodeSelectors; - private final Map workerIsolatedKubeNodeSelectors; + private final Optional> workerIsolatedKubeNodeSelectors; private final Map workerKubeAnnotations; private final String jobImagePullSecret; private final String jobImagePullPolicy; @@ -42,7 +43,7 @@ public WorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), configs.getJobKubeNodeSelectors(), - configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : configs.getJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? Optional.of(configs.getIsolatedJobKubeNodeSelectors()) : Optional.empty(), configs.getJobKubeAnnotations(), configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -74,7 +75,7 @@ public static WorkerConfigs buildSpecWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, - configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : nodeSelectors, + configs.getUseCustomKubeNodeSelector() ? Optional.of(configs.getIsolatedJobKubeNodeSelectors()) : Optional.empty(), annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -106,7 +107,7 @@ public static WorkerConfigs buildCheckWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getCheckJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, - configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : nodeSelectors, + configs.getUseCustomKubeNodeSelector() ? Optional.of(configs.getIsolatedJobKubeNodeSelectors()) : Optional.empty(), annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -138,7 +139,7 @@ public static WorkerConfigs buildDiscoverWorkerConfigs(final Configs configs) { .withMemoryLimit(configs.getJobMainContainerMemoryLimit()), configs.getJobKubeTolerations(), nodeSelectors, - configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : nodeSelectors, + configs.getUseCustomKubeNodeSelector() ? Optional.of(configs.getIsolatedJobKubeNodeSelectors()) : Optional.empty(), annotations, configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -159,7 +160,7 @@ public static WorkerConfigs buildReplicationWorkerConfigs(final Configs configs) .withMemoryLimit(configs.getReplicationOrchestratorMemoryLimit()), configs.getJobKubeTolerations(), configs.getJobKubeNodeSelectors(), - configs.getUseCustomKubeNodeSelector() ? configs.getIsolatedJobKubeNodeSelectors() : configs.getJobKubeNodeSelectors(), + configs.getUseCustomKubeNodeSelector() ? Optional.of(configs.getIsolatedJobKubeNodeSelectors()) : Optional.empty(), configs.getJobKubeAnnotations(), configs.getJobKubeMainContainerImagePullSecret(), configs.getJobKubeMainContainerImagePullPolicy(), @@ -186,7 +187,7 @@ public Map getworkerKubeNodeSelectors() { return workerKubeNodeSelectors; } - public Map getWorkerIsolatedKubeNodeSelectors() { + public Optional> getWorkerIsolatedKubeNodeSelectors() { return workerIsolatedKubeNodeSelectors; } diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java index a9ee7109326f..5dee1d1bbb4d 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java @@ -57,21 +57,26 @@ public class AirbyteIntegrationLauncher implements IntegrationLauncher { private final ResourceRequirements resourceRequirement; private final FeatureFlags featureFlags; - private final boolean useIsolatedNodePool; + /** + * If true, launcher will use a separated isolated pool to run the job. + * + * At this moment, we put custom connector jobs into an isolated pool. + */ + private final boolean useIsolatedPool; public AirbyteIntegrationLauncher(final String jobId, final int attempt, final String imageName, final ProcessFactory processFactory, final ResourceRequirements resourceRequirement, - final boolean useIsolatedNodePool) { + final boolean useIsolatedPool) { this.jobId = jobId; this.attempt = attempt; this.imageName = imageName; this.processFactory = processFactory; this.resourceRequirement = resourceRequirement; this.featureFlags = new EnvVariableFeatureFlags(); - this.useIsolatedNodePool = useIsolatedNodePool; + this.useIsolatedPool = useIsolatedPool; } @Trace(operationName = WORKER_OPERATION_NAME) @@ -84,7 +89,7 @@ public Process spec(final Path jobRoot) throws WorkerException { attempt, jobRoot, imageName, - useIsolatedNodePool, + useIsolatedPool, false, Collections.emptyMap(), null, @@ -105,7 +110,7 @@ public Process check(final Path jobRoot, final String configFilename, final Stri attempt, jobRoot, imageName, - useIsolatedNodePool, + useIsolatedPool, false, ImmutableMap.of(configFilename, configContents), null, @@ -127,7 +132,7 @@ public Process discover(final Path jobRoot, final String configFilename, final S attempt, jobRoot, imageName, - useIsolatedNodePool, + useIsolatedPool, false, ImmutableMap.of(configFilename, configContents), null, @@ -173,7 +178,7 @@ public Process read(final Path jobRoot, attempt, jobRoot, imageName, - useIsolatedNodePool, + useIsolatedPool, false, files, null, @@ -203,7 +208,7 @@ public Process write(final Path jobRoot, attempt, jobRoot, imageName, - useIsolatedNodePool, + useIsolatedPool, true, files, null, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index e0516a162453..a1199ff21679 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -108,7 +108,8 @@ public Process create( final var allLabels = getLabels(jobId, attempt, customLabels); - final var nodeSelectors = usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors() : workerConfigs.getworkerKubeNodeSelectors(); + // If using isolated pool, check workerConfigs has isolated pool set. If not set, fall back to use regular node pool. + final var nodeSelectors = usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); return new KubePodProcess( isOrchestrator, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java index ff1836234b89..65c499f5a114 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/DbtLauncherWorker.java @@ -44,6 +44,8 @@ public DbtLauncherWorker(final UUID connectionId, serverPort, temporalUtils, workerConfigs, + // Custom connector does not use Dbt at this moment, thus this flag for runnning job under + // isolated pool can be set to false. false); } diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java index f078ab41de46..2934740dec91 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java @@ -176,13 +176,17 @@ public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException log.info("Creating " + podName + " for attempt number: " + jobRunConfig.getAttemptId()); killRunningPodsForConnection(); + // custom connectors run in an isolated node pool from airbyte-supported connectors + // to reduce the blast radius of any problems with custom connector code. + final var nodeSelectors = isCustomConnector ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); + try { process.create( allLabels, resourceRequirements, fileMap, portMap, - isCustomConnector ? workerConfigs.getWorkerIsolatedKubeNodeSelectors() : workerConfigs.getworkerKubeNodeSelectors()); + nodeSelectors); } catch (final KubernetesClientException e) { ApmTraceUtils.addExceptionToTrace(e); throw new WorkerException( diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java index dd274e76ef86..38c36fbf52a7 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/NormalizationLauncherWorker.java @@ -45,6 +45,8 @@ public NormalizationLauncherWorker(final UUID connectionId, serverPort, temporalUtils, workerConfigs, + // Normalization process will happen only on a fixed set of connectors, + // thus they are not going to be run under custom connectors. Setting this to false. false); } diff --git a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java index 97bf4af3bb6c..9e628032023b 100644 --- a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java +++ b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java @@ -101,6 +101,8 @@ public Optional runJob() throws Exception { ApmTraceUtils.addTagsToTrace(Map.of(JOB_ID_KEY, jobRunConfig.getJobId(), DESTINATION_DOCKER_IMAGE_KEY, destinationLauncherConfig.getDockerImage(), SOURCE_DOCKER_IMAGE_KEY, sourceLauncherConfig.getDockerImage())); + // At this moment, if either source or destination is from custom connector image, we will put all jobs into isolated pool to run. + boolean useIsolatedPool = sourceLauncherConfig.getIsCustomConnector() || destinationLauncherConfig.getIsCustomConnector(); log.info("Setting up source launcher..."); final IntegrationLauncher sourceLauncher = new AirbyteIntegrationLauncher( sourceLauncherConfig.getJobId(), @@ -108,7 +110,7 @@ public Optional runJob() throws Exception { sourceLauncherConfig.getDockerImage(), processFactory, syncInput.getSourceResourceRequirements(), - sourceLauncherConfig.getIsCustomConnector()); + useIsolatedPool); log.info("Setting up destination launcher..."); final IntegrationLauncher destinationLauncher = new AirbyteIntegrationLauncher( @@ -117,7 +119,7 @@ public Optional runJob() throws Exception { destinationLauncherConfig.getDockerImage(), processFactory, syncInput.getDestinationResourceRequirements(), - destinationLauncherConfig.getIsCustomConnector()); + useIsolatedPool); log.info("Setting up source..."); // reset jobs use an empty source to induce resetting all data in destination. diff --git a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java index 8ee1d21b05b7..116114e061cc 100644 --- a/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java +++ b/airbyte-persistence/job-persistence/src/main/java/io/airbyte/persistence/job/factory/DefaultSyncJobFactory.java @@ -93,7 +93,8 @@ public Long create(final UUID connectionId) { new Version(destinationDefinition.getProtocolVersion()), standardSyncOperations, workspace.getWebhookOperationConfigs(), - sourceDefinition, destinationDefinition) + sourceDefinition, + destinationDefinition) .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-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java index 4223077519a5..9c97cc1349b5 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/DestinationDefinitionsHandler.java @@ -216,7 +216,9 @@ public DestinationDefinitionRead createCustomDestinationDefinition(final CustomD private StandardDestinationDefinition destinationDefinitionFromCreate(final DestinationDefinitionCreate destinationDefCreate) throws IOException { final ConnectorSpecification spec = getSpecForImage( destinationDefCreate.getDockerRepository(), - destinationDefCreate.getDockerImageTag(), true); + destinationDefCreate.getDockerImageTag(), + // Only custom connectors can be created via handlers. + true); final Version airbyteProtocolVersion = AirbyteProtocolVersion.getWithDefault(spec.getProtocolVersion()); diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index 79d6eeb7c6aa..e19d8f9d1f21 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -246,7 +246,11 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceId(final Source final boolean bustActorCatalogCache = discoverSchemaRequestBody.getDisableCache() != null && discoverSchemaRequestBody.getDisableCache(); if (currentCatalog.isEmpty() || bustActorCatalogCache) { final SynchronousResponse persistedCatalogId = - synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, connectorVersion, new Version(sourceDef.getProtocolVersion()), + synchronousSchedulerClient.createDiscoverSchemaJob( + source, + imageName, + connectorVersion, + new Version(sourceDef.getProtocolVersion()), isCustomConnector); final SourceDiscoverSchemaRead discoveredSchema = retrieveDiscoveredSchema(persistedCatalogId); @@ -286,8 +290,13 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceCreate(final So .withSourceDefinitionId(sourceCreate.getSourceDefinitionId()) .withConfiguration(partialConfig) .withWorkspaceId(sourceCreate.getWorkspaceId()); - final SynchronousResponse response = synchronousSchedulerClient.createDiscoverSchemaJob(source, imageName, sourceDef.getDockerImageTag(), - new Version(sourceDef.getProtocolVersion()), isCustomConnector); + final SynchronousResponse response = synchronousSchedulerClient.createDiscoverSchemaJob( + source, + imageName, + sourceDef.getDockerImageTag(), + new Version( + sourceDef.getProtocolVersion()), + isCustomConnector); return retrieveDiscoveredSchema(response); } diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java index 392b792fb0fe..13dce62597f5 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SourceDefinitionsHandler.java @@ -219,7 +219,11 @@ public SourceDefinitionRead createCustomSourceDefinition(final CustomSourceDefin private StandardSourceDefinition sourceDefinitionFromCreate(final SourceDefinitionCreate sourceDefinitionCreate) throws IOException { final ConnectorSpecification spec = - getSpecForImage(sourceDefinitionCreate.getDockerRepository(), sourceDefinitionCreate.getDockerImageTag(), true); + getSpecForImage( + sourceDefinitionCreate.getDockerRepository(), + sourceDefinitionCreate.getDockerImageTag(), + // Only custom connectors can be created via handlers. + true); final Version airbyteProtocolVersion = AirbyteProtocolVersion.getWithDefault(spec.getProtocolVersion()); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index d7daa5bad8c2..4d09b85cf095 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -86,14 +86,14 @@ public Map specJobKubeAnnotations(@Value("${airbyte.worker.spec. } @Singleton - @Named("customNodeSelectors") - public Map customNodeSelectors(@Value("${airbyte.worker.custom.kube.node-selectors}") final String kubeNodeSelectors) { + @Named("isolatedNodeSelectors") + public Map isolatedNodeSelectors(@Value("${airbyte.worker.isolated.kube.node-selectors}") final String kubeNodeSelectors) { return splitKVPairsFromEnvString(kubeNodeSelectors); } @Singleton - @Named("useCustomNodeSelector") - public boolean useCustomNodeSelector(@Value("${airbyte.worker.custom.kube.use-custom-node-selector}") final boolean kubeNodeSelectors) { + @Named("useIsolatedNodeSelector") + public boolean useIsolatedNodeSelector(@Value("${airbyte.worker.isolated.kube.use-isolated-node-selector}") final boolean kubeNodeSelectors) { return kubeNodeSelectors; } @@ -218,8 +218,8 @@ public WorkerConfigs checkWorkerConfigs( @Named("checkResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("checkJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, + @Named("isolatedNodeSelectors") final Map isolatedNodeSelectors, + @Named("useIsolatedNodeSelector") final boolean useIsolatedNodeSelector, @Named("checkJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -233,7 +233,7 @@ public WorkerConfigs checkWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - useCustomNodeSelector ? customNodeSelectors : nodeSelectors, + useIsolatedNodeSelector ? Optional.of(isolatedNodeSelectors) : Optional.empty(), annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -251,8 +251,8 @@ public WorkerConfigs defaultWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, + @Named("isolatedNodeSelectors") final Map isolatedNodeSelectors, + @Named("useIsolatedNodeSelector") final boolean useIsolatedNodeSelector, @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -266,7 +266,7 @@ public WorkerConfigs defaultWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - useCustomNodeSelector ? customNodeSelectors : nodeSelectors, + useIsolatedNodeSelector ? Optional.of(isolatedNodeSelectors) : Optional.empty(), annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -285,8 +285,8 @@ public WorkerConfigs discoverWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("discoverJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, + @Named("isolatedNodeSelectors") final Map isolatedNodeSelectors, + @Named("useIsolatedNodeSelector") final boolean useIsolatedNodeSelector, @Named("discoverJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -300,7 +300,7 @@ public WorkerConfigs discoverWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - useCustomNodeSelector ? customNodeSelectors : nodeSelectors, + useIsolatedNodeSelector ? Optional.of(isolatedNodeSelectors) : Optional.empty(), annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -318,8 +318,8 @@ public WorkerConfigs replicationWorkerConfigs( @Named("replicationResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, + @Named("isolatedNodeSelectors") final Map isolatedNodeSelectors, + @Named("useIsolatedNodeSelector") final boolean useIsolatedNodeSelector, @Named("defaultJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -333,7 +333,7 @@ public WorkerConfigs replicationWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - useCustomNodeSelector ? customNodeSelectors : nodeSelectors, + useIsolatedNodeSelector ? Optional.of(isolatedNodeSelectors) : Optional.empty(), annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, @@ -352,8 +352,8 @@ public WorkerConfigs specWorkerConfigs( @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, final List jobKubeTolerations, @Named("specJobKubeNodeSelectors") final Map nodeSelectors, - @Named("customNodeSelectors") final Map customNodeSelectors, - @Named("useCustomNodeSelector") final boolean useCustomNodeSelector, + @Named("isolatedNodeSelectors") final Map isolatedNodeSelectors, + @Named("useIsolatedNodeSelector") final boolean useIsolatedNodeSelector, @Named("specJobKubeAnnotations") final Map annotations, @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, @@ -367,7 +367,7 @@ public WorkerConfigs specWorkerConfigs( resourceRequirements, jobKubeTolerations, nodeSelectors, - useCustomNodeSelector ? customNodeSelectors : nodeSelectors, + useIsolatedNodeSelector ? Optional.of(isolatedNodeSelectors) : Optional.empty(), annotations, mainContainerImagePullSecret, mainContainerImagePullPolicy, diff --git a/airbyte-workers/src/main/resources/application.yml b/airbyte-workers/src/main/resources/application.yml index 9ec9cd428468..733ae744ebdb 100644 --- a/airbyte-workers/src/main/resources/application.yml +++ b/airbyte-workers/src/main/resources/application.yml @@ -102,10 +102,10 @@ airbyte: memory: limit: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_LIMIT:} request: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_REQUEST:} - custom: + isolated: kube: - use-custom-node-selector: ${USE_CUSTOM_NODE_SELECTOR:false} - node-selectors: ${CUSTOM_JOB_KUBE_NODE_SELECTORS:} + use-isolated-node-selector: ${USE_ISOLATED_NODE_SELECTOR:false} + node-selectors: ${ISOLATED_JOB_KUBE_NODE_SELECTORS:} connection: enabled: ${SHOULD_RUN_CONNECTION_MANAGER_WORKFLOWS:true} discover: From 6fd4f8f975af080b722acbea7ba5e646b01dde8d Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Mon, 5 Dec 2022 12:09:32 -0800 Subject: [PATCH 15/18] style fix --- .../workers/process/AirbyteIntegrationLauncher.java | 2 +- .../io/airbyte/workers/process/KubeProcessFactory.java | 7 +++++-- .../main/java/io/airbyte/workers/sync/LauncherWorker.java | 4 +++- .../container_orchestrator/ReplicationJobOrchestrator.java | 3 ++- .../java/io/airbyte/server/handlers/SchedulerHandler.java | 2 +- 5 files changed, 12 insertions(+), 6 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java index 5dee1d1bbb4d..3d6871786731 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/AirbyteIntegrationLauncher.java @@ -59,7 +59,7 @@ public class AirbyteIntegrationLauncher implements IntegrationLauncher { /** * If true, launcher will use a separated isolated pool to run the job. - * + * * At this moment, we put custom connector jobs into an isolated pool. */ private final boolean useIsolatedPool; diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index a1199ff21679..027480e50dfb 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -108,8 +108,11 @@ public Process create( final var allLabels = getLabels(jobId, attempt, customLabels); - // If using isolated pool, check workerConfigs has isolated pool set. If not set, fall back to use regular node pool. - final var nodeSelectors = usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); + // If using isolated pool, check workerConfigs has isolated pool set. If not set, fall back to use + // regular node pool. + final var nodeSelectors = + usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) + : workerConfigs.getworkerKubeNodeSelectors(); return new KubePodProcess( isOrchestrator, diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java index 2934740dec91..ed2274be290e 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/sync/LauncherWorker.java @@ -178,7 +178,9 @@ public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException // custom connectors run in an isolated node pool from airbyte-supported connectors // to reduce the blast radius of any problems with custom connector code. - final var nodeSelectors = isCustomConnector ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); + final var nodeSelectors = + isCustomConnector ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) + : workerConfigs.getworkerKubeNodeSelectors(); try { process.create( diff --git a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java index 9e628032023b..8fcee891a255 100644 --- a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java +++ b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ReplicationJobOrchestrator.java @@ -101,7 +101,8 @@ public Optional runJob() throws Exception { ApmTraceUtils.addTagsToTrace(Map.of(JOB_ID_KEY, jobRunConfig.getJobId(), DESTINATION_DOCKER_IMAGE_KEY, destinationLauncherConfig.getDockerImage(), SOURCE_DOCKER_IMAGE_KEY, sourceLauncherConfig.getDockerImage())); - // At this moment, if either source or destination is from custom connector image, we will put all jobs into isolated pool to run. + // At this moment, if either source or destination is from custom connector image, we will put all + // jobs into isolated pool to run. boolean useIsolatedPool = sourceLauncherConfig.getIsCustomConnector() || destinationLauncherConfig.getIsCustomConnector(); log.info("Setting up source launcher..."); final IntegrationLauncher sourceLauncher = new AirbyteIntegrationLauncher( diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index e19d8f9d1f21..7bbe6d50514f 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -296,7 +296,7 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceCreate(final So sourceDef.getDockerImageTag(), new Version( sourceDef.getProtocolVersion()), - isCustomConnector); + isCustomConnector); return retrieveDiscoveredSchema(response); } From d23797467a49ab6064e7cbda74222b9955f69591 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 6 Dec 2022 12:03:42 -0800 Subject: [PATCH 16/18] comment fix --- .../workers/process/KubeProcessFactory.java | 5 +++++ .../config/WorkerConfigurationBeanFactory.java | 15 +++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index 027480e50dfb..d28163528fb4 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -110,6 +110,11 @@ public Process create( // If using isolated pool, check workerConfigs has isolated pool set. If not set, fall back to use // regular node pool. + + if (usesIsolatedPool && workerConfigs.getWorkerIsolatedKubeNodeSelectors().isEmpty()) { + throw new WorkerException("Invalid configuration on isolated pool. Isolated kube node " + + "selectors cannot be empty if we want to uses isolated pool"); + } final var nodeSelectors = usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index 4d09b85cf095..d577dc4ef604 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -228,6 +228,9 @@ public WorkerConfigs checkWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { + throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); + } return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -261,6 +264,9 @@ public WorkerConfigs defaultWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { + throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); + } return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -295,6 +301,9 @@ public WorkerConfigs discoverWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { + throw new RuntimeException("Isolated Node selectors is empty while indicating useIsolatedNodeSelector"); + } return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -328,6 +337,9 @@ public WorkerConfigs replicationWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { + throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); + } return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -362,6 +374,9 @@ public WorkerConfigs specWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { + throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); + } return new WorkerConfigs( workerEnvironment, resourceRequirements, From af74bc0af5d8a8ef9c3065113c8f6ef0477ea3b4 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 6 Dec 2022 15:51:55 -0800 Subject: [PATCH 17/18] no checks on kubeprocess --- .../workers/process/KubeProcessFactory.java | 5 ---- .../WorkerConfigurationBeanFactory.java | 26 ++++++++----------- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index d28163528fb4..027480e50dfb 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -110,11 +110,6 @@ public Process create( // If using isolated pool, check workerConfigs has isolated pool set. If not set, fall back to use // regular node pool. - - if (usesIsolatedPool && workerConfigs.getWorkerIsolatedKubeNodeSelectors().isEmpty()) { - throw new WorkerException("Invalid configuration on isolated pool. Isolated kube node " - + "selectors cannot be empty if we want to uses isolated pool"); - } final var nodeSelectors = usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java index d577dc4ef604..aeaf2b044df6 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -210,6 +210,12 @@ public ResourceRequirements replicationResourceRequirements( .withMemoryLimit(memoryLimit); } + void validateIsolatedPoolConfigInitialization(boolean useIsolatedNodeSelector, Map isolatedNodeSelectors) { + if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { + throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); + } + } + @Singleton @Requires(env = WorkerMode.CONTROL_PLANE) @Named("checkWorkerConfigs") @@ -228,9 +234,7 @@ public WorkerConfigs checkWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { - if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { - throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); - } + validateIsolatedPoolConfigInitialization(useIsolatedNodeSelector, isolatedNodeSelectors); return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -264,9 +268,7 @@ public WorkerConfigs defaultWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { - if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { - throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); - } + validateIsolatedPoolConfigInitialization(useIsolatedNodeSelector, isolatedNodeSelectors); return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -301,9 +303,7 @@ public WorkerConfigs discoverWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { - if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { - throw new RuntimeException("Isolated Node selectors is empty while indicating useIsolatedNodeSelector"); - } + validateIsolatedPoolConfigInitialization(useIsolatedNodeSelector, isolatedNodeSelectors); return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -337,9 +337,7 @@ public WorkerConfigs replicationWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { - if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { - throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); - } + validateIsolatedPoolConfigInitialization(useIsolatedNodeSelector, isolatedNodeSelectors); return new WorkerConfigs( workerEnvironment, resourceRequirements, @@ -374,9 +372,7 @@ public WorkerConfigs specWorkerConfigs( @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { - if (useIsolatedNodeSelector && isolatedNodeSelectors.isEmpty()) { - throw new RuntimeException("Isolated Node selectors is empty while useIsolatedNodeSelector is set to true."); - } + validateIsolatedPoolConfigInitialization(useIsolatedNodeSelector, isolatedNodeSelectors); return new WorkerConfigs( workerEnvironment, resourceRequirements, From 219552ba9ba65d41f9e2e2a0018bdbbc23c21b99 Mon Sep 17 00:00:00 2001 From: Xiaohan Song Date: Tue, 6 Dec 2022 15:57:32 -0800 Subject: [PATCH 18/18] rename --- .../java/io/airbyte/workers/process/KubeProcessFactory.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java index 027480e50dfb..ad5b25d7939d 100644 --- a/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java +++ b/airbyte-commons-worker/src/main/java/io/airbyte/workers/process/KubeProcessFactory.java @@ -85,7 +85,7 @@ public Process create( final int attempt, final Path jobRoot, final String imageName, - final boolean usesIsolatedPool, + final boolean isCustomConnector, final boolean usesStdin, final Map files, final String entrypoint, @@ -111,7 +111,7 @@ public Process create( // If using isolated pool, check workerConfigs has isolated pool set. If not set, fall back to use // regular node pool. final var nodeSelectors = - usesIsolatedPool ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) + isCustomConnector ? workerConfigs.getWorkerIsolatedKubeNodeSelectors().orElse(workerConfigs.getworkerKubeNodeSelectors()) : workerConfigs.getworkerKubeNodeSelectors(); return new KubePodProcess(