feat: allow adding more units from env, #9

README.md

@@ -44,8 +44,23 @@ Use `---` to separate multiple units in single `YAML` file
 
 ### 2.2 From Environment Variable
 
+#### Prefix with `MINIT_UNIT_XXXX_`
+
 **Example:**
 
+```dockerfile
+ENV MINIT_UNIT_MAIN_COMMAND="redis-server /etc/redis.conf"
+ENV MINIT_UNIT_MAIN_DIR="/work"
+ENV MINIT_UNIT_MAIN_NAME="main-program"
+ENV MINIT_UNIT_MAIN_GROUP="super-main"
+ENV MINIT_UNIT_MAIN_KIND="cron"
+ENV MINIT_UNIT_MAIN_IMMEDIATE=true
+ENV MINIT_UNIT_MAIN_CRON="* * * * *"
+ENV MINIT_UNIT_MAIN_CHARSET=gbk18030
+```
+
+#### DEPRECATED: `MINIT_MAIN`
+
 ```dockerfile
 ENV MINIT_MAIN="redis-server /etc/redis.conf"
 ENV MINIT_MAIN_DIR="/work"

README.zh.md

@@ -178,15 +178,28 @@ command:
 
 **使用环境变量创建单元**
 
+```dockerfile
+ENV MINIT_UNIT_MAIN_COMMAND="redis-server /etc/redis.conf"
+ENV MINIT_UNIT_MAIN_DIR="/work"
+ENV MINIT_UNIT_MAIN_NAME="main-program"
+ENV MINIT_UNIT_MAIN_GROUP="super-main"
+ENV MINIT_UNIT_MAIN_KIND="cron"
+ENV MINIT_UNIT_MAIN_IMMEDIATE=true
+ENV MINIT_UNIT_MAIN_CRON="* * * * *"
+ENV MINIT_UNIT_MAIN_CHARSET=gbk18030
 ```
-MINIT_MAIN=redis-server /etc/redis.conf
-MINIT_MAIN_DIR=/work
-MINIT_MAIN_NAME=main-program
-MINIT_MAIN_GROUP=super-main
-MINIT_MAIN_KIND=cron
-MINIT_MAIN_CRON="* * * * *"
-MINIT_MAIN_IMMEDIATE=true
-MINIT_MAIN_CHARSET=gbk18030
+
+**已废弃的语法**
+
+```dockerfile
+ENV MINIT_MAIN=redis-server /etc/redis.conf
+ENV MINIT_MAIN_DIR=/work
+ENV MINIT_MAIN_NAME=main-program
+ENV MINIT_MAIN_GROUP=super-main
+ENV MINIT_MAIN_KIND=cron
+ENV MINIT_MAIN_CRON="* * * * *"
+ENV MINIT_MAIN_IMMEDIATE=true
+ENV MINIT_MAIN_CHARSET=gbk18030
 ```
 
 **使用命令行参数创建单元**

go.mod

@@ -3,11 +3,11 @@ module github.com/yankeguo/minit
 go 1.21
 
 require (
-	github.com/yankeguo/rg v1.1.0
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/stretchr/testify v1.8.4
-	golang.org/x/net v0.19.0
-	golang.org/x/sys v0.15.0
+	github.com/yankeguo/rg v1.1.0
+	golang.org/x/net v0.20.0
+	golang.org/x/sys v0.16.0
 	golang.org/x/text v0.14.0
 	gopkg.in/yaml.v3 v3.0.1
 )

go.sum

@@ -6,14 +6,12 @@ github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
 github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
-github.com/yankeguo/rg v1.0.1 h1:Rnca+1JYfuGqPRMIQkuxAoZhhmrPpMFyS5XwLz0U0ds=
-github.com/yankeguo/rg v1.0.1/go.mod h1:tLaoLk8bo/PQld1xGvJvAfCl3K0Nckzh0gsnykFoQYg=
 github.com/yankeguo/rg v1.1.0 h1:kt6RuHzHqHeDctJUcMf9m7J5xaNqGxBin9nQcaveDnU=
 github.com/yankeguo/rg v1.1.0/go.mod h1:OBQueah3CKlMksIbYKNGWJUF20pyy/mipY9NXDXlJ+c=
-golang.org/x/net v0.19.0 h1:zTwKpTd2XuCqf8huc7Fo2iSy+4RHPd10s4KzeTnVr1c=
-golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U=
-golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
-golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo=
+golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
+golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
+golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

pkg/munit/load.go

@@ -56,6 +56,138 @@ func LoadArgs(args []string) (unit Unit, ok bool, err error) {
 	return
 }
 
+const (
+	UnitPrefix = "MINIT_UNIT_"
+)
+
+// DetectEnvInfixes detects infixes from environment variables
+func DetectEnvInfixes() (infixes []string) {
+	_infixes := map[string]struct{}{}
+
+	for _, item := range os.Environ() {
+		splits := strings.SplitN(item, "=", 2)
+		if len(splits) != 2 {
+			continue
+		}
+
+		key := splits[0]
+		if !strings.HasPrefix(key, UnitPrefix) {
+			continue
+		}
+		key = strings.TrimPrefix(key, UnitPrefix)
+		if strings.HasSuffix(key, "_COMMAND") {
+			key = strings.TrimSuffix(key, "_COMMAND")
+			if key == "" {
+				continue
+			}
+			_infixes[key] = struct{}{}
+		} else if strings.HasSuffix(key, "_FILES") {
+			key = strings.TrimSuffix(key, "_FILES")
+			if key == "" {
+				continue
+			}
+			if os.Getenv(UnitPrefix+key+"_KIND") == KindRender {
+				_infixes[key] = struct{}{}
+			}
+		}
+	}
+
+	for infix := range _infixes {
+		infixes = append(infixes, infix)
+	}
+
+	return
+}
+
+// LoadFromEnvWithInfix loads unit from environment variables with infix
+// e.g. MINIT_MAIN_KIND, MINIT_MAIN_NAME, MINIT_MAIN_COMMAND
+func LoadFromEnvWithInfix(infix string) (unit Unit, ok bool, err error) {
+	// kind
+	unit.Kind = os.Getenv(UnitPrefix + infix + "_KIND")
+
+	if unit.Kind == "" {
+		unit.Kind = KindDaemon
+	}
+
+	switch unit.Kind {
+	case KindDaemon, KindOnce, KindCron, KindRender:
+	default:
+		err = errors.New("unsupported $" + UnitPrefix + infix + "_KIND: " + unit.Kind)
+		return
+	}
+
+	// name, group, count
+	unit.Name = os.Getenv(UnitPrefix + infix + "_NAME")
+	if unit.Name == "" {
+		unit.Name = "env-" + strings.ToLower(infix)
+	}
+	unit.Group = os.Getenv(UnitPrefix + infix + "_GROUP")
+	unit.Count, _ = strconv.Atoi(os.Getenv("MINIT_" + infix + "_COUNT"))
+
+	// command, dir, shell, charset
+	switch unit.Kind {
+	case KindDaemon, KindOnce, KindCron:
+		cmd := os.Getenv(UnitPrefix + infix + "_COMMAND")
+
+		if unit.Command, err = shellquote.Split(cmd); err != nil {
+			return
+		}
+
+		if len(unit.Command) == 0 {
+			err = errors.New("missing environment variable $MINIT_" + infix + "_COMMAND")
+			return
+		}
+
+		unit.Dir = os.Getenv(UnitPrefix + infix + "_DIR")
+		unit.Shell = os.Getenv(UnitPrefix + infix + "_SHELL")
+		unit.Charset = os.Getenv(UnitPrefix + infix + "_CHARSET")
+
+		for _, item := range strings.Split(os.Getenv(UnitPrefix+infix+"_ENV"), ";") {
+			item = strings.TrimSpace(item)
+			splits := strings.SplitN(item, "=", 2)
+			if len(splits) == 2 {
+				if unit.Env == nil {
+					unit.Env = make(map[string]string)
+				}
+				unit.Env[splits[0]] = splits[1]
+			}
+		}
+	}
+
+	// cron, immediate
+	if unit.Kind == KindCron {
+		unit.Cron = os.Getenv(UnitPrefix + infix + "_CRON")
+
+		if unit.Cron == "" {
+			err = errors.New("missing environment variable $" + UnitPrefix + infix + "_CRON while $" + UnitPrefix + infix + "_KIND is 'cron'")
+			return
+		}
+
+		unit.Immediate, _ = strconv.ParseBool(os.Getenv(UnitPrefix + infix + "_IMMEDIATE"))
+	}
+
+	// raw, files
+	if unit.Kind == KindRender {
+		unit.Raw, _ = strconv.ParseBool(os.Getenv(UnitPrefix + infix + "_RAW"))
+
+		for _, item := range strings.Split(os.Getenv(UnitPrefix+infix+"_FILES"), ";") {
+			item = strings.TrimSpace(item)
+			if item != "" {
+				unit.Files = append(unit.Files, item)
+			}
+		}
+
+		if len(unit.Files) == 0 {
+			err = errors.New("missing environment variable $" + UnitPrefix + infix + "_FILES while $" + UnitPrefix + infix + "_KIND is 'render'")
+			return
+		}
+	}
+
+	ok = true
+
+	return
+}
+
 func LoadEnv() (unit Unit, ok bool, err error) {
 	cmd := strings.TrimSpace(os.Getenv("MINIT_MAIN"))
 	if cmd == "" {

pkg/munit/loader.go

@@ -60,13 +60,30 @@ func (ld *Loader) Load(opts LoadOptions) (output []Unit, skipped []Unit, err err
 		}
 	}
 	if opts.Env {
-		var unit Unit
-		var ok bool
-		if unit, ok, err = LoadEnv(); err != nil {
-			return
+		{
+			// legacy minit main
+			var (
+				unit Unit
+				ok   bool
+			)
+			if unit, ok, err = LoadEnv(); err != nil {
+				return
+			} else if ok {
+				units = append(units, unit)
+			}
 		}
-		if ok {
-			units = append(units, unit)
+
+		for _, infix := range DetectEnvInfixes() {
+			var (
+				unit Unit
+				ok   bool
+			)
+			if unit, ok, err = LoadFromEnvWithInfix(infix); err != nil {
+				return
+			}
+			if ok {
+				units = append(units, unit)
+			}
 		}
 	}
 

pkg/munit/loader_test.go

@@ -1,9 +1,10 @@
 package munit
 
 import (
-	"github.com/stretchr/testify/require"
 	"os"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewLoader(t *testing.T) {
@@ -34,3 +35,34 @@ func TestDupOrMakeMap(t *testing.T) {
 	require.Equal(t, "d", m1a["c"])
 	require.Equal(t, "", m1b["c"])
 }
+
+func TestLoaderWithNewEnv(t *testing.T) {
+	os.Unsetenv("MINIT_ENABLE")
+	os.Unsetenv("MINIT_DISABLE")
+	os.Unsetenv("MINIT_MAIN_NAME")
+	os.Unsetenv("MINIT_MAIN_KIND")
+	os.Unsetenv("MINIT_MAIN_ONCE")
+	os.Setenv("MINIT_MAIN", "legacy main")
+	os.Setenv("MINIT_UNIT_CACHE_COMMAND", "redis-server")
+	os.Setenv("MINIT_UNIT_INITIAL_COMMAND", "touch /tmp/initial")
+	os.Setenv("MINIT_UNIT_INITIAL_NAME", "job-initial")
+	os.Setenv("MINIT_UNIT_INITIAL_ENV", "ZAA=ZBB;ZCC=ZDD")
+	os.Setenv("MINIT_UNIT_INITIAL_KIND", "once")
+
+	l := NewLoader()
+	units, _, err := l.Load(LoadOptions{Env: true})
+	require.NoError(t, err)
+	require.Len(t, units, 3)
+
+	require.Equal(t, "env-main", units[0].Name)
+	require.Equal(t, []string{"legacy", "main"}, units[0].Command)
+	require.Equal(t, KindDaemon, units[0].Kind)
+
+	require.Equal(t, "env-cache", units[1].Name)
+	require.Equal(t, []string{"redis-server"}, units[1].Command)
+	require.Equal(t, KindDaemon, units[1].Kind)
+
+	require.Equal(t, "job-initial", units[2].Name)
+	require.Equal(t, []string{"touch", "/tmp/initial"}, units[2].Command)
+	require.Equal(t, KindOnce, units[2].Kind)
+}