jpictl/pkg/zones/env.go

package zones

import (
	"bytes"
	"fmt"
	"io"
	"strings"
)

// 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,
	}
}

// Zones returns the list of Zone IDs
func (m *Env) Zones() []int {
	var zones []int

	m.ForEachZone(func(z *Zone) bool {
		zones = append(zones, z.ID)
		return false
	})

	return zones
}

// WriteTo generates environment variables for shell scripts
func (m *Env) WriteTo(w io.Writer) (int64, error) {
	var buf bytes.Buffer

	m.writeEnvVarInts(&buf, m.Zones(), "ZONES")
	m.ForEachZone(func(z *Zone) bool {
		m.writeEnvZone(&buf, z)
		return false
	})

	return buf.WriteTo(w)
}

func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
	zoneID := z.ID

	// ZONE{zoneID}
	m.writeEnvVar(w, genEnvZoneNodes(z), "ZONE%v", zoneID)

	// ZONE{zoneID}_NAME
	m.writeEnvVar(w, z.Name, "ZONE%v_%s", zoneID, "NAME")

	// ZONE{zoneID}_GW
	gateways, _ := z.GatewayIDs()
	m.writeEnvVarInts(w, gateways, "ZONE%v_%s", zoneID, "GW")
}

func (m *Env) writeEnvVarInts(w io.Writer, value []int, name string, args ...any) {
	var s string

	if n := len(value); n > 0 {
		var buf bytes.Buffer

		for i, v := range value {
			if i != 0 {
				_, _ = fmt.Fprint(&buf, " ")
			}
			_, _ = fmt.Fprintf(&buf, "%v", v)
		}

		s = buf.String()
	}

	m.writeEnvVar(w, s, name, args...)
}

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...)
	}

	if name != "" {
		value = strings.TrimSpace(value)

		_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
	}
}

func genEnvZoneNodes(z *Zone) string {
	if n := z.Len(); n > 0 {
		s := make([]string, 0, n)

		z.ForEachMachine(func(p *Machine) bool {
			s = append(s, p.Name)
			return false
		})

		return strings.Join(s, " ")
	}
	return ""
}