initial commit

pkg/watermill/azsb/azbus.go

@@ -0,0 +1,80 @@
+package azsb
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/rs/zerolog"
+)
+
+type AzBus struct {
+	client      *azservicebus.Client
+	logger      zerolog.Logger
+	closed      bool
+	closedMutex sync.RWMutex
+}
+
+type Config struct {
+	ConnectionString   string
+	UseManagedIdentity bool
+	Namespace          string
+}
+
+// NewAzBus creates a new Azure Service Bus publisher and subscriber
+func NewAzBus(cfg Config, logger zerolog.Logger) (message.Subscriber, message.Publisher, error) {
+	var client *azservicebus.Client
+	var err error
+
+	if cfg.UseManagedIdentity {
+		// Use managed identity
+		if cfg.Namespace == "" {
+			return nil, nil, fmt.Errorf("azure service bus namespace is required when using managed identity")
+		}
+
+		credential, credErr := azidentity.NewDefaultAzureCredential(nil)
+		if credErr != nil {
+			return nil, nil, fmt.Errorf("failed to create azure credential: %w", credErr)
+		}
+
+		namespace := cfg.Namespace
+		client, err = azservicebus.NewClient(namespace, credential, nil)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to create azure service bus client: %w", err)
+		}
+	} else {
+		// Use connection string
+		if cfg.ConnectionString == "" {
+			return nil, nil, fmt.Errorf("azure service bus connection string is not configured")
+		}
+
+		client, err = azservicebus.NewClientFromConnectionString(cfg.ConnectionString, nil)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to create azure service bus client: %w", err)
+		}
+	}
+
+	azb := &AzBus{client: client, logger: logger, closed: false, closedMutex: sync.RWMutex{}}
+
+	return azb, azb, nil
+}
+
+func (a *AzBus) Close() error {
+	if a.closed {
+		return nil
+	}
+
+	if a.client != nil {
+		if err := a.client.Close(context.Background()); err != nil {
+			a.logger.Error().Err(err).Msg("failed to close azure service bus client")
+			return err
+		}
+	}
+
+	a.closed = true
+	a.logger.Info().Msg("azure service bus publisher closed")
+	return nil
+}

pkg/watermill/azsb/publisher.go

@@ -0,0 +1,65 @@
+package azsb
+
+import (
+	"context"
+	"fmt"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+func (a *AzBus) Publish(topic string, messages ...*message.Message) error {
+	if a.closed {
+		return fmt.Errorf("publisher is closed")
+	}
+
+	if len(messages) == 0 {
+		return nil
+	}
+
+	sender, err := a.client.NewSender(topic, nil)
+	if err != nil {
+		return fmt.Errorf("failed to create sender for topic %s: %w", topic, err)
+	}
+	defer sender.Close(context.Background())
+
+	sbMessages := new(azservicebus.MessageBatch)
+	for _, msg := range messages {
+		sbMsg := &azservicebus.Message{
+			Body: msg.Payload,
+		}
+
+		// Copy metadata as application properties
+		if msg.Metadata != nil {
+			sbMsg.ApplicationProperties = make(map[string]interface{})
+			for key, value := range msg.Metadata {
+				sbMsg.ApplicationProperties[key] = value
+			}
+		}
+
+		// Set message ID if available
+		if msg.UUID != "" {
+			sbMsg.MessageID = &msg.UUID
+		}
+
+		err = sbMessages.AddMessage(sbMsg, nil)
+		if err != nil {
+			return err
+		}
+	}
+
+	if err = sender.SendMessageBatch(context.Background(), sbMessages, nil); err != nil {
+		a.logger.Error().
+			Err(err).
+			Str("topic", topic).
+			Int32("message_count", sbMessages.NumMessages()).
+			Msg("failed to send messages to azure service bus")
+		return fmt.Errorf("failed to send messages: %w", err)
+	}
+
+	a.logger.Debug().
+		Str("topic", topic).
+		Int32("message_count", sbMessages.NumMessages()).
+		Msg("published messages to azure service bus")
+
+	return nil
+}

pkg/watermill/azsb/subscriber.go

@@ -0,0 +1,125 @@
+package azsb
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+	"github.com/ThreeDotsLabs/watermill/message"
+)
+
+func (a *AzBus) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
+	a.closedMutex.RLock()
+	if a.closed {
+		a.closedMutex.RUnlock()
+		return nil, fmt.Errorf("subscriber is closed")
+	}
+	a.closedMutex.RUnlock()
+
+	// Create receiver for the subscription
+	// In Azure Service Bus, you need to create a subscription for a topic before subscribing
+	// The subscription name should match what was created in Azure Service Bus
+	// Default: use topic name with "-subscription" suffix
+	// You should create the subscription in Azure Service Bus beforehand or make this configurable
+	subscriptionName := topic + "-subscription"
+
+	receiver, err := a.client.NewReceiverForSubscription(topic, subscriptionName, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create receiver for topic %s subscription %s: %w. Note: Subscription must be created in Azure Service Bus first", topic, subscriptionName, err)
+	}
+
+	messages := make(chan *message.Message, 100)
+
+	go func() {
+		defer close(messages)
+		defer receiver.Close(context.Background())
+
+		for {
+			select {
+			case <-ctx.Done():
+				a.logger.Info().Str("topic", topic).Msg("subscription context cancelled")
+				return
+			default:
+				// Check if closed
+				a.closedMutex.RLock()
+				if a.closed {
+					a.closedMutex.RUnlock()
+					return
+				}
+				a.closedMutex.RUnlock()
+
+				// Receive messages
+				messages2, err := receiver.ReceiveMessages(ctx, 1, nil)
+				if err != nil {
+					if ctx.Err() != nil {
+						return
+					}
+					a.logger.Error().
+						Err(err).
+						Str("topic", topic).
+						Msg("failed to receive messages from azure service bus")
+					continue
+				}
+
+				for _, sbMsg := range messages2 {
+					watermillMsg := a.convertToWatermillMessage(sbMsg)
+
+					select {
+					case messages <- watermillMsg:
+						// Message sent successfully
+						// Complete the message
+						if err := receiver.CompleteMessage(ctx, sbMsg, nil); err != nil {
+							a.logger.Error().
+								Err(err).
+								Str("message_id", watermillMsg.UUID).
+								Msg("failed to complete message")
+						}
+					case <-ctx.Done():
+						// Context cancelled, abandon the message
+						if err := receiver.AbandonMessage(ctx, sbMsg, nil); err != nil {
+							a.logger.Error().
+								Err(err).
+								Str("message_id", watermillMsg.UUID).
+								Msg("failed to abandon message")
+						}
+						return
+					}
+				}
+			}
+		}
+	}()
+
+	a.logger.Info().
+		Str("topic", topic).
+		Str("subscription", subscriptionName).
+		Msg("started subscribing to azure service bus")
+
+	return messages, nil
+}
+
+func (a *AzBus) convertToWatermillMessage(sbMsg *azservicebus.ReceivedMessage) *message.Message {
+	msg := message.NewMessage("", sbMsg.Body)
+
+	// Set message ID
+	if sbMsg.MessageID != "=" {
+		msg.UUID = sbMsg.MessageID
+	}
+
+	// Copy application properties to metadata
+	if sbMsg.ApplicationProperties != nil {
+		msg.Metadata = make(message.Metadata)
+		for key, value := range sbMsg.ApplicationProperties {
+			if strValue, ok := value.(string); ok {
+				msg.Metadata[key] = strValue
+			} else {
+				// Convert non-string values to string
+				if jsonValue, err := json.Marshal(value); err == nil {
+					msg.Metadata[key] = string(jsonValue)
+				}
+			}
+		}
+	}
+
+	return msg
+}