pkg/cluster/env.go

@@ -4,12 +4,14 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"sort"
 	"strings"
 )
 
 // Env is a shell environment factory for this cluster
 type Env struct {
 	ZoneIterator
+	RegionIterator
 
 	cephFSID string
 	export   bool
@@ -23,9 +25,10 @@ func (m *Cluster) Env(export bool) (*Env, error) {
 	}
 
 	env := &Env{
-		ZoneIterator: m,
-		cephFSID:     fsid.String(),
-		export:       export,
+		ZoneIterator:   m,
+		RegionIterator: m,
+		cephFSID:       fsid.String(),
+		export:         export,
 	}
 
 	return env, nil
@@ -43,6 +46,22 @@ func (m *Env) Zones() []int {
 	return zones
 }
 
+// Regions returns the list of primary regions
+func (m *Env) Regions() []string {
+	var regions []string
+
+	m.ForEachRegion(func(r *Region) bool {
+		if r.Cluster != nil {
+			regions = append(regions, r.Name)
+		}
+
+		return false
+	})
+
+	sort.Strings(regions)
+	return regions
+}
+
 // WriteTo generates environment variables for shell scripts
 func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
@@ -51,6 +70,7 @@ func (m *Env) WriteTo(w io.Writer) (int64, error) {
 		m.writeEnvVar(&buf, m.cephFSID, "FSID")
 	}
 
+	m.writeEnvVarStrings(&buf, m.Regions(), "REGIONS")
 	m.writeEnvVarInts(&buf, m.Zones(), "ZONES")
 
 	m.ForEachZone(func(z *Zone) bool {
@@ -74,6 +94,9 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	gateways, _ := z.GatewayIDs()
 	m.writeEnvVarInts(w, gateways, "ZONE%v_%s", zoneID, "GW")
 
+	// ZONE{zoneID}_REGION
+	m.writeEnvVar(w, genEnvZoneRegion(z), "ZONE%v_%s", zoneID, "REGION")
+
 	// Ceph
 	monitors := z.GetCephMonitors()
 	// MON{zoneID}_NAME
@@ -85,22 +108,29 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 }
 
 func (m *Env) writeEnvVarInts(w io.Writer, value []int, name string, args ...any) {
-	var s string
+	var buf bytes.Buffer
 
-	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)
+	for _, v := range value {
+		if buf.Len() > 0 {
+			_, _ = fmt.Fprint(&buf, " ")
 		}
-
-		s = buf.String()
+		_, _ = fmt.Fprintf(&buf, "%v", v)
 	}
 
-	m.writeEnvVar(w, s, name, args...)
+	m.writeEnvVar(w, buf.String(), name, args...)
+}
+
+func (m *Env) writeEnvVarStrings(w io.Writer, value []string, name string, args ...any) {
+	var buf bytes.Buffer
+
+	for _, v := range value {
+		if buf.Len() > 0 {
+			_, _ = fmt.Fprint(&buf, " ")
+		}
+		_, _ = fmt.Fprintf(&buf, "%s", v)
+	}
+
+	m.writeEnvVar(w, buf.String(), name, args...)
 }
 
 func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
@@ -117,7 +147,11 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	if name != "" {
 		value = strings.TrimSpace(value)
 
-		_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
+		if value == "" {
+			_, _ = fmt.Fprintf(w, "%s%s=\n", prefix, name)
+		} else {
+			_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
+		}
 	}
 }
 
@@ -135,6 +169,21 @@ func genEnvZoneNodes(z *Zone) string {
 	return ""
 }
 
+func genEnvZoneRegion(z *Zone) string {
+	var region string
+
+	z.ForEachRegion(func(r *Region) bool {
+		if r.Cluster != nil {
+			region = r.Name
+			return true
+		}
+
+		return false
+	})
+
+	return region
+}
+
 func genEnvZoneCephMonNames(m Machines) string {
 	var buf strings.Builder
 	m.ForEachMachine(func(p *Machine) bool {