From 8043209629f1958748aaf68dbe5daceae735e853 Mon Sep 17 00:00:00 2001
From: Naya Chen <nchen@confluent.io>
Date: Thu, 6 Jan 2022 16:30:11 -0800
Subject: [PATCH] Add command args to container creation (#6)

formatting

formatting
---
 apis/druid/v1alpha1/druid_types.go            |  9 ++++++++-
 .../crds/druid.apache.org_druids.yaml         |  4 ++++
 controllers/druid/handler.go                  | 19 ++++++++++++++++++-
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/apis/druid/v1alpha1/druid_types.go b/apis/druid/v1alpha1/druid_types.go
index 99cd6b47..c887139a 100644
--- a/apis/druid/v1alpha1/druid_types.go
+++ b/apis/druid/v1alpha1/druid_types.go
@@ -53,9 +53,16 @@ type DruidSpec struct {
 	// +optional
 	DeleteOrphanPvc bool `json:"deleteOrphanPvc"`
 
-	// Required: path to druid start script to be run on container start
+	// Required: Command to be run on container start
 	StartScript string `json:"startScript"`
 
+	// Optional: bash/sh entry arg. Set startScript to `sh` or `bash` to customize entryArg
+	// For example, the container can run `sh -c "${EntryArg} && ${DruidScript} {nodeType}"`
+	EntryArg string `json:"entryArg,omitempty"`
+
+	// Optional: Customized druid shell script path. If not set, the default would be "bin/run-druid.sh"
+	DruidScript string `json:"druidScript,omitempty"`
+
 	// Required here or at nodeSpec level
 	Image string `json:"image,omitempty"`
 
diff --git a/chart/templates/crds/druid.apache.org_druids.yaml b/chart/templates/crds/druid.apache.org_druids.yaml
index 38ff27fd..6aa7f3dc 100644
--- a/chart/templates/crds/druid.apache.org_druids.yaml
+++ b/chart/templates/crds/druid.apache.org_druids.yaml
@@ -3700,6 +3700,10 @@ spec:
                 type: array
               startScript:
                 type: string
+              entryArg:
+                type: string
+              druidScript:
+                type: string
               startUpProbes:
                 properties:
                   exec:
diff --git a/controllers/druid/handler.go b/controllers/druid/handler.go
index 6de60af4..9811ad2d 100644
--- a/controllers/druid/handler.go
+++ b/controllers/druid/handler.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"regexp"
 	"sort"
+	"strings"
 
 	autoscalev2beta2 "k8s.io/api/autoscaling/v2beta2"
 	networkingv1 "k8s.io/api/networking/v1"
@@ -1138,6 +1139,21 @@ func getVolume(nodeSpec *v1alpha1.DruidNodeSpec, m *v1alpha1.Druid, nodeSpecUniq
 	return volumesHolder
 }
 
+func getCommand(nodeSpec *v1alpha1.DruidNodeSpec, m *v1alpha1.Druid) []string {
+  if (m.Spec.StartScript != "" && m.Spec.EntryArg != "") {
+      return []string{m.Spec.StartScript}
+  }
+  return []string{firstNonEmptyStr(m.Spec.StartScript, "bin/run-druid.sh"), nodeSpec.NodeType}
+}
+
+func getEntryArg(nodeSpec *v1alpha1.DruidNodeSpec, m *v1alpha1.Druid) []string {
+  if (m.Spec.EntryArg != "") {
+      bashCommands := strings.Join([]string{m.Spec.EntryArg, "&&", firstNonEmptyStr(m.Spec.DruidScript, "bin/run-druid.sh"), nodeSpec.NodeType}, " ")
+      return []string{"-c", bashCommands}
+  }
+  return nil
+}
+
 func getEnv(nodeSpec *v1alpha1.DruidNodeSpec, m *v1alpha1.Druid, configMapSHA string) []v1.EnvVar {
 	envHolder := firstNonNilValue(nodeSpec.Env, m.Spec.Env).([]v1.EnvVar)
 	// enables to do the trick to force redeployment in case of configmap changes.
@@ -1312,7 +1328,8 @@ func makePodSpec(nodeSpec *v1alpha1.DruidNodeSpec, m *v1alpha1.Druid, nodeSpecUn
 			{
 				Image:           firstNonEmptyStr(nodeSpec.Image, m.Spec.Image),
 				Name:            fmt.Sprintf("%s", nodeSpecUniqueStr),
-				Command:         []string{firstNonEmptyStr(m.Spec.StartScript, "bin/run-druid.sh"), nodeSpec.NodeType},
+				Command:         getCommand(nodeSpec, m),
+				Args:            getEntryArg(nodeSpec, m),
 				ImagePullPolicy: v1.PullPolicy(firstNonEmptyStr(string(nodeSpec.ImagePullPolicy), string(m.Spec.ImagePullPolicy))),
 				Ports:           nodeSpec.Ports,
 				Resources:       nodeSpec.Resources,