Updating to use context

Author: Michael Freeman

client.go

@@ -32,13 +32,13 @@ type Client struct {
 type messageHandler func(context.Context, jetstream.Msg) error
 
 // Connect establishes a connection to NATS and sets up JetStream.
-func (c *Client) Connect() error {
+func (c *Client) Connect(ctx context.Context) error {
 	if err := c.validateAndPrepare(); err != nil {
 		return err
 	}
 
 	connManager := NewConnectionManager(c.Config, c.logger, c.natsConnector, c.jetStreamCreator)
-	if err := connManager.Connect(); err != nil {
+	if err := connManager.Connect(ctx); err != nil {
 		c.logger.Errorf("failed to connect to NATS server at %v: %v", c.Config.Server, err)
 		return err
 	}

connection_manager.go

@@ -82,7 +82,7 @@ func NewConnectionManager(
 }
 
 // Connect establishes a connection to NATS and sets up JetStream.
-func (cm *ConnectionManager) Connect() error {
+func (cm *ConnectionManager) Connect(ctx context.Context) error {
 	cm.logger.Logf("Connecting to NATS server at %v", cm.config.Server)
 
 	opts := []nats.Option{nats.Name("GoFr NATS JetStreamClient")}

interfaces.go

@@ -42,7 +42,7 @@ type JetStreamClient interface {
 
 // ConnectionManagerInterface represents the main Client connection.
 type ConnectionManagerInterface interface {
-	Connect() error
+	Connect(ctx context.Context) error
 	Close(ctx context.Context)
 	Publish(ctx context.Context, subject string, message []byte, metrics Metrics) error
 	Health() datasource.Health

pubsub_wrapper.go

@@ -44,14 +44,14 @@ func (w *PubSubWrapper) Health() datasource.Health {
 }
 
 // Connect establishes a connection to NATS.
-func (w *PubSubWrapper) Connect() error {
+func (w *PubSubWrapper) Connect(ctx context.Context) error {
 	if w.Client.connManager != nil && w.Client.connManager.Health().Status == datasource.StatusUp {
 		w.Client.logger.Log("NATS connection already established")
 
 		return nil
 	}
 
-	err := w.Client.Connect()
+	err := w.Client.Connect(ctx)
 	if err != nil {
 		w.Client.logger.Errorf("PubSubWrapper: Error connecting to NATS: %v", err)
 		return err