Add support for private image registry

airbytehq · Oct 20, 2023 · 1b06309 · 1b06309
1 parent fa6fb0d
commit 1b06309
Show file tree

Hide file tree

Showing 9 changed files with 23 additions and 3 deletions.
diff --git a/...mons-with-dependencies/src/main/java/io/airbyte/commons/workers/config/WorkerConfigs.java b/...mons-with-dependencies/src/main/java/io/airbyte/commons/workers/config/WorkerConfigs.java
@@ -26,6 +26,7 @@ public class WorkerConfigs {
   private final Map<String, String> workerKubeAnnotations;
   private final Map<String, String> workerKubeLabels;
   private final List<String> jobImagePullSecrets;
+  private final String jobImageRegistry;
   private final String jobImagePullPolicy;
   private final String sidecarImagePullPolicy;
   private final String jobSocatImage;
@@ -41,6 +42,7 @@ public WorkerConfigs(final WorkerEnvironment workerEnvironment,
                        final Map<String, String> workerKubeAnnotations,
                        final Map<String, String> workerKubeLabels,
                        final List<String> jobImagePullSecrets,
+                       final String jobImageRegistry,
                        final String jobImagePullPolicy,
                        final String sidecarImagePullPolicy,
                        final String jobSocatImage,
@@ -55,6 +57,7 @@ public WorkerConfigs(final WorkerEnvironment workerEnvironment,
     this.workerKubeAnnotations = workerKubeAnnotations;
     this.workerKubeLabels = workerKubeLabels;
     this.jobImagePullSecrets = jobImagePullSecrets;
+    this.jobImageRegistry = jobImageRegistry;
     this.jobImagePullPolicy = jobImagePullPolicy;
     this.sidecarImagePullPolicy = sidecarImagePullPolicy;
     this.jobSocatImage = jobSocatImage;
@@ -78,6 +81,7 @@ public WorkerConfigs(final Configs configs) {
         configs.getJobKubeAnnotations(),
         configs.getJobKubeLabels(),
         configs.getJobKubeMainContainerImagePullSecrets(),
+        null,
         configs.getJobKubeMainContainerImagePullPolicy(),
         configs.getJobKubeSidecarContainerImagePullPolicy(),
         configs.getJobKubeSocatImage(),
@@ -118,6 +122,10 @@ public List<String> getJobImagePullSecrets() {
     return jobImagePullSecrets;
   }
 
+  public String getJobImageRegistry() {
+    return jobImageRegistry;
+  }
+
   public String getJobImagePullPolicy() {
     return jobImagePullPolicy;
   }

diff --git a/...h-dependencies/src/main/java/io/airbyte/commons/workers/config/WorkerConfigsProvider.java b/...h-dependencies/src/main/java/io/airbyte/commons/workers/config/WorkerConfigsProvider.java
@@ -170,6 +170,7 @@ record WorkerConfigsDefaults(WorkerEnvironment workerEnvironment,
                                @Value("${airbyte.worker.isolated.kube.node-selectors}") String isolatedNodeSelectors,
                                @Value("${airbyte.worker.isolated.kube.use-custom-node-selector}") boolean useCustomNodeSelector,
                                @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") List<String> mainContainerImagePullSecret,
+                               @Value("${airbyte.worker.job.kube.main.container.image-registry}") String mainContainerImageRegistry,
                                @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") String mainContainerImagePullPolicy,
                                @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") String sidecarContainerImagePullPolicy,
                                @Value("${airbyte.worker.job.kube.images.socat}") String socatImage,
@@ -249,6 +250,7 @@ private WorkerConfigs getConfig(final KubeResourceKey key) {
         annotations,
         splitKVPairsFromEnvString(kubeResourceConfig.getLabels()),
         workerConfigsDefaults.mainContainerImagePullSecret(),
+        workerConfigsDefaults.mainContainerImageRegistry(),
         workerConfigsDefaults.mainContainerImagePullPolicy(),
         workerConfigsDefaults.sidecarContainerImagePullPolicy(),
         workerConfigsDefaults.socatImage(),

diff --git a/airbyte-commons-with-dependencies/src/test/resources/application-config-test.yaml b/airbyte-commons-with-dependencies/src/test/resources/application-config-test.yaml
@@ -60,6 +60,7 @@ airbyte:
           socat: ${UNUSED_VALUE:}
         main:
           container:
+            image-registry: ${UNUSED_VALUE:}
             image-pull-policy: ${UNUSED_VALUE:}
             image-pull-secret: ${UNUSED_VALUE:}
         sidecar:

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
@@ -135,7 +135,10 @@ public Process create(final ResourceType resourceType,
           rebasePath(jobRoot).toString(), // rebases the job root on the job data mount
           "--log-driver",
           "none");
-      final String containerName = ProcessFactory.createProcessName(imageName, jobType, jobId, attempt, DOCKER_NAME_LEN_LIMIT);
+      final WorkerConfigs workerConfigs = workerConfigsProvider.getConfig(resourceType);     
+      final String registry = workerConfigs.getJobImageRegistry();
+      final String image = null == registry || imageName.startsWith(registry) ? imageName : registry + "/" + imageName;
+      final String containerName = ProcessFactory.createProcessName(image, jobType, jobId, attempt, DOCKER_NAME_LEN_LIMIT);
       final ResourceRequirements resourceRequirements = connectorResourceRequirements.main();
       LOGGER.info("Creating docker container = {} with resources {} and allowedHosts {}", containerName, resourceRequirements, allowedHosts);
       cmd.add("--name");
@@ -158,7 +161,6 @@ public Process create(final ResourceType resourceType,
         cmd.add(String.format("%s:%s", localMountSource, LOCAL_MOUNT_DESTINATION));
       }
 
-      final WorkerConfigs workerConfigs = workerConfigsProvider.getConfig(resourceType);
       final Map<String, String> allEnvMap = MoreMaps.merge(jobMetadata, workerConfigs.getEnvMap());
       for (final Map.Entry<String, String> envEntry : allEnvMap.entrySet()) {
         cmd.add("-e");

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
@@ -153,6 +153,9 @@ public Process create(
 
       final Context featureFlagContext = createFeatureFlagContext(connectionId, workspaceId, imageName);
 
+      final String registry = workerConfigs.getJobImageRegistry();
+      final String image = null == registry || imageName.startsWith(registry) ? imageName : registry + "/" + imageName;
+
       return new KubePodProcess(
           processRunnerHost,
           fabricClient,
@@ -162,7 +165,7 @@ public Process create(
           namespace,
           serviceAccount,
           schedulerName.isBlank() ? null : schedulerName,
-          imageName,
+          image,
           workerConfigs.getJobImagePullPolicy(),
           workerConfigs.getSidecarImagePullPolicy(),
           stdoutLocalPort,

diff --git a/airbyte-container-orchestrator/src/main/resources/application.yml b/airbyte-container-orchestrator/src/main/resources/application.yml
@@ -52,6 +52,7 @@ airbyte:
           socat: ${JOB_KUBE_SOCAT_IMAGE:`alpine/socat:1.7.4.3-r0`}
         main:
           container:
+            image-registry: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_REGISTRY:}
             image-pull-policy: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_POLICY:IfNotPresent}
             image-pull-secret: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_SECRET:}
         sidecar:

diff --git a/airbyte-container-orchestrator/src/test/resources/application-test.yml b/airbyte-container-orchestrator/src/test/resources/application-test.yml
@@ -37,6 +37,7 @@ airbyte:
           socat: ${JOB_KUBE_SOCAT_IMAGE:`alpine/socat:1.7.4.3-r0`}
         main:
           container:
+            image-registry: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_REGISTRY:}
             image-pull-policy: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_POLICY:IfNotPresent}
             image-pull-secret: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_SECRET:}
         sidecar:

diff --git a/airbyte-server/src/main/resources/application.yml b/airbyte-server/src/main/resources/application.yml
@@ -222,6 +222,7 @@ airbyte:
           socat: ${JOB_KUBE_SOCAT_IMAGE:`alpine/socat:1.7.4.3-r0`}
         main:
           container:
+            image-registry: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_REGISTRY:}
             image-pull-policy: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_POLICY:IfNotPresent}
             image-pull-secret: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_SECRET:}
         sidecar:

diff --git a/airbyte-workers/src/main/resources/application.yml b/airbyte-workers/src/main/resources/application.yml
@@ -197,6 +197,7 @@ airbyte:
           socat: ${JOB_KUBE_SOCAT_IMAGE:`alpine/socat:1.7.4.3-r0`}
         main:
           container:
+            image-registry: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_REGISTRY:}
             image-pull-policy: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_POLICY:IfNotPresent}
             image-pull-secret: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_SECRET:}
         namespace: ${JOB_KUBE_NAMESPACE:default}