discovery service

api/pkg/messaging/internal/natsb/broker.go

@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"os"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -20,6 +21,7 @@ type natsSubscriotions = map[string]*TopicSubscription
 
 type NatsBroker struct {
 	nc        *nats.Conn
+	js        nats.JetStreamContext
 	logger    *zap.Logger
 	topicSubs natsSubscriotions
 	mu        sync.Mutex
@@ -78,23 +80,46 @@ func loadEnv(settings *nc.Settings, l *zap.Logger) (*envConfig, error) {
 
 func NewNatsBroker(logger mlogger.Logger, settings *nc.Settings) (*NatsBroker, error) {
 	l := logger.Named("broker")
-	// Helper function to get environment variables
-	cfg, err := loadEnv(settings, l)
-	if err != nil {
-		return nil, err
+	var err error
+	var cfg *envConfig
+	var natsURL string
+	if settings != nil && strings.TrimSpace(settings.URLEnv) != "" {
+		urlVal := strings.TrimSpace(os.Getenv(settings.URLEnv))
+		if urlVal != "" {
+			natsURL = urlVal
+		}
 	}
+	if natsURL == "" {
+		// Helper function to get environment variables
+		cfg, err = loadEnv(settings, l)
+		if err != nil {
+			return nil, err
+		}
 
-	u := &url.URL{
-		Scheme: "nats",
-		Host:   net.JoinHostPort(cfg.Host, strconv.Itoa(cfg.Port)),
+		u := &url.URL{
+			Scheme: "nats",
+			Host:   net.JoinHostPort(cfg.Host, strconv.Itoa(cfg.Port)),
+		}
+		natsURL = u.String()
 	}
-	natsURL := u.String()
 
 	opts := []nats.Option{
 		nats.Name(settings.NATSName),
 		nats.MaxReconnects(settings.MaxReconnects),
 		nats.ReconnectWait(time.Duration(settings.ReconnectWait) * time.Second),
-		nats.UserInfo(cfg.User, cfg.Password),
+	}
+	if cfg != nil {
+		opts = append(opts, nats.UserInfo(cfg.User, cfg.Password))
+	} else if settings != nil {
+		userEnv := strings.TrimSpace(settings.UsernameEnv)
+		passEnv := strings.TrimSpace(settings.PasswordEnv)
+		if userEnv != "" && passEnv != "" {
+			user := strings.TrimSpace(os.Getenv(userEnv))
+			pass := strings.TrimSpace(os.Getenv(passEnv))
+			if user != "" || pass != "" {
+				opts = append(opts, nats.UserInfo(user, pass))
+			}
+		}
 	}
 
 	res := &NatsBroker{
@@ -106,8 +131,18 @@ func NewNatsBroker(logger mlogger.Logger, settings *nc.Settings) (*NatsBroker, e
 		l.Error("Failed to connect to NATS", zap.String("url", natsURL), zap.Error(err))
 		return nil, err
 	}
+	if res.js, err = res.nc.JetStream(); err != nil {
+		l.Warn("Failed to initialise JetStream context", zap.Error(err))
+	}
 
 	logger.Info("Connected to NATS", zap.String("broker", settings.NATSName),
-		zap.String("url", fmt.Sprintf("nats://%s@%s", cfg.User, net.JoinHostPort(cfg.Host, strconv.Itoa(cfg.Port)))))
+		zap.String("url", natsURL))
 	return res, nil
 }
+
+func (b *NatsBroker) JetStream() nats.JetStreamContext {
+	if b == nil {
+		return nil
+	}
+	return b.js
+}