Merge pull request 'env: introduce Env() factory and `jpictl env -e` to export variables' (#2)

Reviewed-on: #2

cmd/jpictl/env.go

@@ -16,10 +16,19 @@ var envCmd = &cobra.Command{
 			return err
 		}
 
-		return m.WriteEnv(os.Stdout)
+		_, err = m.Env(*envExport).WriteTo(os.Stdout)
+		return err
 	},
 }
 
+// Command Flags
+var (
+	envExport *bool
+)
+
 func init() {
 	rootCmd.AddCommand(envCmd)
+
+	envExport = envCmd.PersistentFlags().BoolP("export", "e", false,
+		"export generated variables")
 }

pkg/zones/env.go

@@ -7,8 +7,23 @@ import (
 	"strings"
 )
 
-// WriteEnv generates environment variables for shell scripts
-func (m *Zones) WriteEnv(w io.Writer) error {
+// Env is a shell environment factory for this cluster
+type Env struct {
+	ZoneIterator
+
+	export bool
+}
+
+// Env returns a shell environment factory
+func (m *Zones) Env(export bool) *Env {
+	return &Env{
+		ZoneIterator: m,
+		export:       export,
+	}
+}
+
+// WriteTo generates environment variables for shell scripts
+func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
 
 	m.writeEnvVarFn(&buf, genEnvZones, "ZONES")
@@ -17,11 +32,10 @@ func (m *Zones) WriteEnv(w io.Writer) error {
 		return false
 	})
 
-	_, err := buf.WriteTo(w)
-	return err
+	return buf.WriteTo(w)
 }
 
-func (m *Zones) writeEnvZone(w io.Writer, z *Zone) {
+func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	zoneID := z.ID
 
 	// ZONE{zoneID}
@@ -42,7 +56,7 @@ func (m *Zones) writeEnvZone(w io.Writer, z *Zone) {
 	}
 }
 
-func (m *Zones) writeEnvVarFn(w io.Writer, fn func(*Zones) string, name string, args ...any) {
+func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) {
 	var value string
 
 	if fn != nil {
@@ -52,7 +66,13 @@ func (m *Zones) writeEnvVarFn(w io.Writer, fn func(*Zones) string, name string,
 	m.writeEnvVar(w, value, name, args...)
 }
 
-func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
+func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
+	var prefix string
+
+	if m.export {
+		prefix = "export "
+	}
+
 	if len(args) > 0 {
 		name = fmt.Sprintf(name, args...)
 	}
@@ -60,15 +80,17 @@ func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	if name != "" {
 		value = strings.TrimSpace(value)
 
-		_, _ = fmt.Fprintf(w, "%s=%q\n", name, value)
+		_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
 	}
 }
 
-func genEnvZones(m *Zones) string {
-	s := make([]string, 0, len(m.Zones))
-	for _, z := range m.Zones {
+func genEnvZones(m *Env) string {
+	var s []string
+
+	m.ForEachZone(func(z *Zone) bool {
 		s = append(s, fmt.Sprintf("%v", z.ID))
-	}
+		return false
+	})
 
 	return strings.Join(s, " ")
 }