You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

114 lines
2.0 KiB

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,
}
}
// 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")
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
gatewayID := getRingZeroGatewayID(z)
if gatewayID > 0 {
m.writeEnvVar(w, fmt.Sprintf("%v", gatewayID), "ZONE%v_%s", zoneID, "GW")
}
}
func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) {
var value string
if fn != nil {
value = fn(m)
}
m.writeEnvVar(w, value, 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 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, " ")
}
func genEnvZoneNodes(z *Zone) string {
s := make([]string, 0, len(z.Machines))
for _, p := range z.Machines {
s = append(s, p.Name)
}
return strings.Join(s, " ")
}
func getRingZeroGatewayID(z *Zone) int {
var gatewayID int
z.ForEachMachine(func(p *Machine) bool {
if p.IsGateway() {
gatewayID = p.ID
}
return gatewayID != 0
})
return gatewayID
}