support init container in k8s deploy

Signed-off-by: Patrick Zhao <[email protected]>
koderover · Sep 19, 2024 · 817f42c · 817f42c
1 parent 4038996
commit 817f42c
Show file tree

Hide file tree

Showing 9 changed files with 180 additions and 48 deletions.
diff --git a/pkg/microservice/aslan/core/common/repository/models/service.go b/pkg/microservice/aslan/core/common/repository/models/service.go
@@ -151,10 +151,11 @@ type ImagePathSpec struct {
 
 // Container ...
 type Container struct {
-	Name      string         `bson:"name"                          json:"name"`
-	Image     string         `bson:"image"                         json:"image"`
-	ImageName string         `bson:"image_name,omitempty"          json:"image_name,omitempty"`
-	ImagePath *ImagePathSpec `bson:"image_path,omitempty"          json:"imagePath,omitempty"`
+	Name      string                `bson:"name"                          json:"name"`
+	Type      setting.ContainerType `bson:"type"                          json:"type"`
+	Image     string                `bson:"image"                         json:"image"`
+	ImageName string                `bson:"image_name,omitempty"          json:"image_name,omitempty"`
+	ImagePath *ImagePathSpec        `bson:"image_path,omitempty"          json:"imagePath,omitempty"`
 }
 
 // ServiceTmplPipeResp ...router

diff --git a/pkg/microservice/aslan/core/common/service/kube/render.go b/pkg/microservice/aslan/core/common/service/kube/render.go
@@ -154,6 +154,12 @@ func ReplaceWorkloadImages(rawYaml string, images []*commonmodels.Container) (st
 					deployment.Spec.Template.Spec.Containers[i].Image = image.Image
 				}
 			}
+			for i, container := range deployment.Spec.Template.Spec.InitContainers {
+				containerName := container.Name
+				if image, ok := imageMap[containerName]; ok {
+					deployment.Spec.Template.Spec.InitContainers[i].Image = image.Image
+				}
+			}
 
 			yamlStr, err = resourceToYaml(deployment)
 			if err != nil {
@@ -174,6 +180,12 @@ func ReplaceWorkloadImages(rawYaml string, images []*commonmodels.Container) (st
 					statefulSet.Spec.Template.Spec.Containers[i].Image = image.Image
 				}
 			}
+			for i, container := range statefulSet.Spec.Template.Spec.InitContainers {
+				containerName := container.Name
+				if image, ok := imageMap[containerName]; ok {
+					statefulSet.Spec.Template.Spec.InitContainers[i].Image = image.Image
+				}
+			}
 			yamlStr, err = resourceToYaml(statefulSet)
 			if err != nil {
 				return "", nil, err
@@ -189,6 +201,12 @@ func ReplaceWorkloadImages(rawYaml string, images []*commonmodels.Container) (st
 					job.Spec.Template.Spec.Containers[i].Image = image.Image
 				}
 			}
+			for i, container := range job.Spec.Template.Spec.InitContainers {
+				containerName := container.Name
+				if image, ok := imageMap[containerName]; ok {
+					job.Spec.Template.Spec.InitContainers[i].Image = image.Image
+				}
+			}
 			yamlStr, err = resourceToYaml(job)
 			if err != nil {
 				return "", nil, err
@@ -209,6 +227,12 @@ func ReplaceWorkloadImages(rawYaml string, images []*commonmodels.Container) (st
 						cronJob.Spec.JobTemplate.Spec.Template.Spec.Containers[i].Image = image.Image
 					}
 				}
+				for i, val := range cronJob.Spec.JobTemplate.Spec.Template.Spec.InitContainers {
+					containerName := val.Name
+					if image, ok := imageMap[containerName]; ok {
+						cronJob.Spec.JobTemplate.Spec.Template.Spec.InitContainers[i].Image = image.Image
+					}
+				}
 				yamlStr, err = resourceToYaml(cronJob)
 				if err != nil {
 					return "", nil, err
@@ -228,6 +252,12 @@ func ReplaceWorkloadImages(rawYaml string, images []*commonmodels.Container) (st
 						cronJob.Spec.JobTemplate.Spec.Template.Spec.Containers[i].Image = image.Image
 					}
 				}
+				for i, val := range cronJob.Spec.JobTemplate.Spec.Template.Spec.InitContainers {
+					containerName := val.Name
+					if image, ok := imageMap[containerName]; ok {
+						cronJob.Spec.JobTemplate.Spec.Template.Spec.InitContainers[i].Image = image.Image
+					}
+				}
 				yamlStr, err = resourceToYaml(cronJob)
 				if err != nil {
 					return "", nil, err

diff --git a/pkg/microservice/aslan/core/common/service/workflowcontroller/jobcontroller/job_deploy.go b/pkg/microservice/aslan/core/common/service/workflowcontroller/jobcontroller/job_deploy.go
@@ -338,6 +338,24 @@ L:
 				break L
 			}
 		}
+
+		for _, container := range deploy.Spec.Template.Spec.InitContainers {
+			if container.Name == serviceModule.ServiceModule {
+				err = updater.UpdateDeploymentInitImage(deploy.Namespace, deploy.Name, serviceModule.ServiceModule, serviceModule.Image, c.kubeClient)
+				if err != nil {
+					return fmt.Errorf("failed to update container image in %s/deployments/%s/%s: %v", env.Namespace, deploy.Name, container.Name, err)
+				}
+				c.jobTaskSpec.ReplaceResources = append(c.jobTaskSpec.ReplaceResources, commonmodels.Resource{
+					Kind:      setting.Deployment,
+					Container: container.Name,
+					Origin:    container.Image,
+					Name:      deploy.Name,
+				})
+				replaced = true
+				c.jobTaskSpec.RelatedPodLabels = append(c.jobTaskSpec.RelatedPodLabels, deploy.Spec.Template.Labels)
+				break L
+			}
+		}
 	}
 Loop:
 	for _, sts := range statefulSets {
@@ -358,6 +376,23 @@ Loop:
 				break Loop
 			}
 		}
+		for _, container := range sts.Spec.Template.Spec.InitContainers {
+			if container.Name == serviceModule.ServiceModule {
+				err = updater.UpdateStatefulSetInitImage(sts.Namespace, sts.Name, serviceModule.ServiceModule, serviceModule.Image, c.kubeClient)
+				if err != nil {
+					return fmt.Errorf("failed to update container image in %s/statefulsets/%s/%s: %v", env.Namespace, sts.Name, container.Name, err)
+				}
+				c.jobTaskSpec.ReplaceResources = append(c.jobTaskSpec.ReplaceResources, commonmodels.Resource{
+					Kind:      setting.StatefulSet,
+					Container: container.Name,
+					Origin:    container.Image,
+					Name:      sts.Name,
+				})
+				replaced = true
+				c.jobTaskSpec.RelatedPodLabels = append(c.jobTaskSpec.RelatedPodLabels, sts.Spec.Template.Labels)
+				break Loop
+			}
+		}
 	}
 CronLoop:
 	for _, cron := range cronJobs {
@@ -378,6 +413,23 @@ CronLoop:
 				break CronLoop
 			}
 		}
+		for _, container := range cron.Spec.JobTemplate.Spec.Template.Spec.InitContainers {
+			if container.Name == serviceModule.ServiceModule {
+				err = updater.UpdateCronJobInitImage(cron.Namespace, cron.Name, serviceModule.ServiceModule, serviceModule.Image, c.kubeClient, false)
+				if err != nil {
+					return fmt.Errorf("failed to update container image in %s/cronJob/%s/%s: %v", env.Namespace, cron.Name, container.Name, err)
+				}
+				c.jobTaskSpec.ReplaceResources = append(c.jobTaskSpec.ReplaceResources, commonmodels.Resource{
+					Kind:      setting.CronJob,
+					Container: container.Name,
+					Origin:    container.Image,
+					Name:      cron.Name,
+				})
+				replaced = true
+				c.jobTaskSpec.RelatedPodLabels = append(c.jobTaskSpec.RelatedPodLabels, cron.Spec.JobTemplate.Spec.Template.Labels)
+				break CronLoop
+			}
+		}
 	}
 BetaCronLoop:
 	for _, cron := range betaCronJobs {
@@ -398,6 +450,23 @@ BetaCronLoop:
 				break BetaCronLoop
 			}
 		}
+		for _, container := range cron.Spec.JobTemplate.Spec.Template.Spec.InitContainers {
+			if container.Name == serviceModule.ServiceModule {
+				err = updater.UpdateCronJobInitImage(cron.Namespace, cron.Name, serviceModule.ServiceModule, serviceModule.Image, c.kubeClient, true)
+				if err != nil {
+					return fmt.Errorf("failed to update container image in %s/cronJobBeta/%s/%s: %v", env.Namespace, cron.Name, container.Name, err)
+				}
+				c.jobTaskSpec.ReplaceResources = append(c.jobTaskSpec.ReplaceResources, commonmodels.Resource{
+					Kind:      setting.CronJob,
+					Container: container.Name,
+					Origin:    container.Image,
+					Name:      cron.Name,
+				})
+				replaced = true
+				c.jobTaskSpec.RelatedPodLabels = append(c.jobTaskSpec.RelatedPodLabels, cron.Spec.JobTemplate.Spec.Template.Labels)
+				break BetaCronLoop
+			}
+		}
 	}
 
 	if !replaced {

diff --git a/pkg/microservice/aslan/core/common/util/service.go b/pkg/microservice/aslan/core/common/util/service.go
@@ -184,30 +184,19 @@ func getContainers(data string) ([]*commonmodels.Container, error) {
 		return containers, fmt.Errorf("unmarshal yaml data error: %v", err)
 	}
 
-	for _, val := range res.Spec.Template.Spec.Containers {
-
-		if _, ok := val["name"]; !ok {
-			return containers, errors.New("yaml file missing name key")
-		}
-
-		nameStr, ok := val["name"].(string)
-		if !ok {
-			return containers, errors.New("error name value")
-		}
-
-		if _, ok := val["image"]; !ok {
-			return containers, errors.New("yaml file missing image key")
+	for _, val := range res.Spec.Template.Spec.InitContainers {
+		container, err := parseContainer(val, setting.ContainerTypeInit)
+		if err != nil {
+			return containers, err
 		}
 
-		imageStr, ok := val["image"].(string)
-		if !ok {
-			return containers, errors.New("error image value")
-		}
+		containers = append(containers, container)
+	}
 
-		container := &commonmodels.Container{
-			Name:      nameStr,
-			Image:     imageStr,
-			ImageName: ExtractImageName(imageStr),
+	for _, val := range res.Spec.Template.Spec.Containers {
+		container, err := parseContainer(val, setting.ContainerTypeNormal)
+		if err != nil {
+			return containers, err
 		}
 
 		containers = append(containers, container)
@@ -224,30 +213,19 @@ func getCronJobContainers(data string) ([]*commonmodels.Container, error) {
 		return containers, fmt.Errorf("unmarshal yaml data error: %v", err)
 	}
 
-	for _, val := range res.Spec.Template.Spec.Template.Spec.Containers {
-
-		if _, ok := val["name"]; !ok {
-			return containers, errors.New("yaml file missing name key")
-		}
-
-		nameStr, ok := val["name"].(string)
-		if !ok {
-			return containers, errors.New("error name value")
-		}
-
-		if _, ok := val["image"]; !ok {
-			return containers, errors.New("yaml file missing image key")
+	for _, val := range res.Spec.Template.Spec.Template.Spec.InitContainers {
+		container, err := parseContainer(val, setting.ContainerTypeInit)
+		if err != nil {
+			return containers, err
 		}
 
-		imageStr, ok := val["image"].(string)
-		if !ok {
-			return containers, errors.New("error image value")
-		}
+		containers = append(containers, container)
+	}
 
-		container := &commonmodels.Container{
-			Name:      nameStr,
-			Image:     imageStr,
-			ImageName: ExtractImageName(imageStr),
+	for _, val := range res.Spec.Template.Spec.Template.Spec.Containers {
+		container, err := parseContainer(val, setting.ContainerTypeNormal)
+		if err != nil {
+			return containers, err
 		}
 
 		containers = append(containers, container)
@@ -504,3 +482,31 @@ func getValuesByPath(paths map[string]string, flatMap map[string]interface{}) ma
 	}
 	return ret
 }
+
+func parseContainer(val map[string]interface{}, containerType setting.ContainerType) (*commonmodels.Container, error) {
+	if _, ok := val["name"]; !ok {
+		return nil, errors.New("yaml file missing name key")
+	}
+
+	nameStr, ok := val["name"].(string)
+	if !ok {
+		return nil, errors.New("error name value")
+	}
+
+	if _, ok := val["image"]; !ok {
+		return nil, errors.New("yaml file missing image key")
+	}
+
+	imageStr, ok := val["image"].(string)
+	if !ok {
+		return nil, errors.New("error image value")
+	}
+
+	container := &commonmodels.Container{
+		Name:      nameStr,
+		Type:      containerType,
+		Image:     imageStr,
+		ImageName: ExtractImageName(imageStr),
+	}
+	return container, nil
+}
diff --git a/pkg/setting/consts.go b/pkg/setting/consts.go
@@ -916,3 +916,10 @@ const (
 	SQLExecStatusFailed  SQLExecStatus = "failed"
 	SQLExecStatusNotExec SQLExecStatus = "not_exec"
 )
+
+type ContainerType string
+
+const (
+	ContainerTypeInit   ContainerType = "init"
+	ContainerTypeNormal ContainerType = ""
+)
diff --git a/pkg/tool/kube/updater/cronjob.go b/pkg/tool/kube/updater/cronjob.go
@@ -72,6 +72,11 @@ func UpdateCronJobImage(ns, name, container, image string, cl client.Client, ver
 	return PatchCronJob(ns, name, patchBytes, cl, versionLessThan121)
 }
 
+func UpdateCronJobInitImage(ns, name, container, image string, cl client.Client, versionLessThan121 bool) error {
+	patchBytes := []byte(fmt.Sprintf(`{"spec":{"jobTemplate":{"spec":{"template":{"spec":{"initContainers":[{"name":"%s","image":"%s"}]}}}}}}`, container, image))
+	return PatchCronJob(ns, name, patchBytes, cl, versionLessThan121)
+}
+
 func SuspendCronJob(ns, name string, cl client.Client, versionLessThan121 bool) error {
 	patchBytes := []byte(`{"spec":{"suspend":true}}`)
 	return PatchCronJob(ns, name, patchBytes, cl, versionLessThan121)

diff --git a/pkg/tool/kube/updater/deployment.go b/pkg/tool/kube/updater/deployment.go
@@ -90,6 +90,12 @@ func UpdateDeploymentImage(ns, name, container, image string, cl client.Client)
 	return PatchDeployment(ns, name, patchBytes, cl)
 }
 
+func UpdateDeploymentInitImage(ns, name, container, image string, cl client.Client) error {
+	patchBytes := []byte(fmt.Sprintf(`{"spec":{"template":{"spec":{"initContainers":[{"name":"%s","image":"%s"}]}}}}`, container, image))
+
+	return PatchDeployment(ns, name, patchBytes, cl)
+}
+
 func ScaleDeployment(ns, name string, replicas int, cl client.Client) error {
 	patchBytes := []byte(fmt.Sprintf(`{"spec":{"replicas": %d}}`, replicas))
 	return PatchDeployment(ns, name, patchBytes, cl)

diff --git a/pkg/tool/kube/updater/statefulset.go b/pkg/tool/kube/updater/statefulset.go
@@ -74,6 +74,12 @@ func UpdateStatefulSetImage(ns, name, container, image string, cl client.Client)
 	return PatchStatefulSet(ns, name, patchBytes, cl)
 }
 
+func UpdateStatefulSetInitImage(ns, name, container, image string, cl client.Client) error {
+	patchBytes := []byte(fmt.Sprintf(`{"spec":{"template":{"spec":{"initContainers":[{"name":"%s","image":"%s"}]}}}}`, container, image))
+
+	return PatchStatefulSet(ns, name, patchBytes, cl)
+}
+
 func ScaleStatefulSet(ns, name string, replicas int, cl client.Client) error {
 	patchBytes := []byte(fmt.Sprintf(`{"spec":{"replicas": %d}}`, replicas))
 	return PatchStatefulSet(ns, name, patchBytes, cl)

diff --git a/pkg/types/workload.go b/pkg/types/workload.go
@@ -54,7 +54,8 @@ type KubeResource struct {
 	Spec struct {
 		Template struct {
 			Spec struct {
-				Containers []map[string]interface{} `yaml:"containers"`
+				InitContainers []map[string]interface{} `yaml:"initContainers"`
+				Containers     []map[string]interface{} `yaml:"containers"`
 			} `yaml:"spec"`
 		} `yaml:"template"`
 	} `yaml:"spec"`
@@ -71,7 +72,8 @@ type CronjobResource struct {
 			Spec struct {
 				Template struct {
 					Spec struct {
-						Containers []map[string]interface{} `yaml:"containers"`
+						InitContainers []map[string]interface{} `yaml:"initContainers"`
+						Containers     []map[string]interface{} `yaml:"containers"`
 					} `yaml:"spec"`
 				} `yaml:"template"`
 			} `yaml:"spec"`