Adding single queue multi topic code

Signed-off-by: Christopher Hein <[email protected]>

Author: christopherhein

pkg/config/types.go

@@ -17,6 +17,8 @@ type Config struct {
 	Region           string
 	Kubeconfig       string
 	MasterURL        string
+	QueueURL         string
+	QueueARN         string
 	AWSSession       *session.Session
 	AWSClientset     awsclient.ServiceoperatorV1alpha1Interface
 	KubeClientset    kubernetes.Interface

pkg/queue/queue.go

@@ -1,74 +1,19 @@
 package queue
 
 import (
+	"context"
 	"encoding/json"
+	"os"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/sns"
 	"github.com/aws/aws-sdk-go/service/sqs"
 	awsclient "github.com/awslabs/aws-service-operator/pkg/client/clientset/versioned/typed/service-operator.aws/v1alpha1"
 	"github.com/awslabs/aws-service-operator/pkg/config"
-	"github.com/awslabs/aws-service-operator/pkg/helpers"
-	"os"
-	"strings"
+	"github.com/awslabs/aws-service-operator/pkg/queuemanager"
 )
 
-// HandlerFunc allows you to define a custom function for when a message is stored
-type HandlerFunc func(config *config.Config, msg *MessageBody) error
-
-// HandleMessage will stub the handler for processing messages
-func (f HandlerFunc) HandleMessage(config *config.Config, msg *MessageBody) error {
-	return f(config, msg)
-}
-
-// Handler allows a custom function to be passed
-type Handler interface {
-	HandleMessage(config *config.Config, msg *MessageBody) error
-}
-
-// ParseMessage will take the message attribute and make it readable
-func (m *MessageBody) ParseMessage() error {
-	m.Updatable = false
-	resp := make(map[string]string)
-	items := strings.Split(m.Message, "\n")
-	for _, item := range items {
-		x := strings.Split(item, "=")
-		key := x[0]
-		if key != "" {
-			s := x[1]
-			s = s[1 : len(s)-1]
-			resp[key] = s
-		}
-	}
-	m.ParsedMessage = resp
-
-	var resourceProperties ResourceProperties
-	if resp["ResourceProperties"] != "null" {
-		err := json.Unmarshal([]byte(resp["ResourceProperties"]), &resourceProperties)
-		if err != nil {
-			return err
-		}
-		m.ResourceProperties = resourceProperties
-		for _, tag := range resourceProperties.Tags {
-			switch tag.Key {
-			case "Namespace":
-				m.Namespace = tag.Value
-			case "ResourceName":
-				m.ResourceName = tag.Value
-			}
-		}
-		if m.Namespace != "" && m.ResourceName != "" {
-			m.Updatable = true
-		}
-	}
-	return nil
-}
-
-// IsComplete returns a simple status instead of the raw CFT resp
-func (m *MessageBody) IsComplete() bool {
-	return helpers.IsStackComplete(m.ParsedMessage["ResourceStatus"], true)
-}
-
 // New will initialize the Queue object for watching
 func New(config *config.Config, awsclientset awsclient.ServiceoperatorV1alpha1Interface, timeout int) *Queue {
 	return &Queue{
@@ -78,21 +23,11 @@ func New(config *config.Config, awsclientset awsclient.ServiceoperatorV1alpha1In
 	}
 }
 
-// Register will create all the affilate resources
-func (q *Queue) Register(name string, obj interface{}) (topicARN string, queueURL string, queueAttrs map[string]*string, subARN string) {
-	config := q.config
-	logger := config.Logger
-	keyname := keyName(config.ClusterName, name)
-	svc := sns.New(config.AWSSession)
-	topicInputs := sns.CreateTopicInput{Name: aws.String(keyname)}
-	output, err := svc.CreateTopic(&topicInputs)
-	if err != nil {
-		logger.Errorf("Error creating SNS Topic with error '%s'", err.Error())
-	}
-	topicARN = *output.TopicArn
-	logger.Infof("Created sns topic '%s'", topicARN)
-
+// RegisterQueue wkll create the Queue so it is accessible to SNS.
+func RegisterQueue(config *config.Config, name string) (queueURL, queueARN string, manager *queuemanager.QueueManager, err error) {
+	manager = queuemanager.New()
 	sqsSvc := sqs.New(config.AWSSession)
+	keyname := keyName(config.ClusterName, name)
 	queueInputs := sqs.CreateQueueInput{
 		QueueName: aws.String(keyname),
 		Attributes: map[string]*string{
@@ -102,14 +37,10 @@ func (q *Queue) Register(name string, obj interface{}) (topicARN string, queueUR
 	}
 	sqsOutput, err := sqsSvc.CreateQueue(&queueInputs)
 	if err != nil {
-		logger.Errorf("Error creating SQS Queue with error '%s'", err.Error())
+		return "", "", manager, err
 	}
 	queueURL = *sqsOutput.QueueUrl
-	logger.Infof("Created sqs queue '%s'", queueURL)
-	q.queueURL = queueURL
 
-	// Get Queue Information
-	// This will later happen with the CRD cft subscription
 	queueQueryInputs := sqs.GetQueueAttributesInput{
 		QueueUrl: aws.String(queueURL),
 		AttributeNames: []*string{
@@ -118,29 +49,46 @@ func (q *Queue) Register(name string, obj interface{}) (topicARN string, queueUR
 	}
 	sqsQueueOutput, err := sqsSvc.GetQueueAttributes(&queueQueryInputs)
 	if err != nil {
-		logger.Errorf("Error getting SQS Information with error '%s'", err.Error())
+		return "", "", manager, err
 	}
-	queueAttrs = sqsQueueOutput.Attributes
+	queueARN = *sqsQueueOutput.Attributes["QueueArn"]
+	return queueURL, queueARN, manager, nil
+}
 
-	policy := newPolicy(*queueAttrs["QueueArn"], keyname, topicARN)
-	policyb, err := json.Marshal(policy)
+// Subscribe will listen to the global queue and distribute messages
+func Subscribe(config *config.Config, manager *queuemanager.QueueManager, ctx context.Context) {
+	logger := config.Logger
+	sess, err := session.NewSession(&aws.Config{
+		Region: aws.String(config.Region),
+	})
 	if err != nil {
-		logger.WithError(err).Error("error encoding policy to json")
+		logger.WithError(err).Error("error creating AWS session")
+		os.Exit(1)
 	}
-	addPermInput := &sqs.SetQueueAttributesInput{
-		QueueUrl: aws.String(queueURL),
-		Attributes: map[string]*string{
-			"Policy": aws.String(string(policyb)),
-		}}
-	_, err = sqsSvc.SetQueueAttributes(addPermInput)
+
+	svc := sqs.New(sess)
+
+	process(config, svc, manager, ctx)
+}
+
+// Register will create all the affilate resources
+func (q *Queue) Register(name string) (topicARN string, subARN string) {
+	config := q.config
+	logger := config.Logger
+	keyname := keyName(config.ClusterName, name)
+	svc := sns.New(config.AWSSession)
+	topicInputs := sns.CreateTopicInput{Name: aws.String(keyname)}
+	output, err := svc.CreateTopic(&topicInputs)
 	if err != nil {
-		logger.WithError(err).Error("error setting queue policy")
+		logger.Errorf("Error creating SNS Topic with error '%s'", err.Error())
 	}
+	topicARN = *output.TopicArn
+	logger.Infof("Created sns topic '%s'", topicARN)
 
 	// Move to somewhere else.
 	subInput := sns.SubscribeInput{
 		TopicArn: aws.String(topicARN),
-		Endpoint: aws.String(*queueAttrs["QueueArn"]),
+		Endpoint: aws.String(config.QueueARN),
 		Protocol: aws.String("sqs"),
 	}
 	subOutput, err := svc.Subscribe(&subInput)
@@ -153,84 +101,96 @@ func (q *Queue) Register(name string, obj interface{}) (topicARN string, queueUR
 	return
 }
 
-// StartWatch will start a go routine to watch for new events to store in etcd
-func (q *Queue) StartWatch(h Handler, stopCh <-chan struct{}) {
-	config := q.config
-	logger := config.Logger
-	sess, err := session.NewSession(&aws.Config{
-		Region: aws.String(config.Region),
-	})
+func SetQueuePolicy(config *config.Config, manager *queuemanager.QueueManager) error {
+	sqsSvc := sqs.New(config.AWSSession)
+	topicARNs := manager.Keys()
+
+	policy := newPolicy(config.QueueARN, config.ClusterName, topicARNs)
+	policyb, err := json.Marshal(policy)
 	if err != nil {
-		logger.WithError(err).Infof("error creating AWS session")
-		os.Exit(1)
+		return err
 	}
 
-	svc := sqs.New(sess)
-
-	process(q, svc, h, stopCh)
+	addPermInput := &sqs.SetQueueAttributesInput{
+		QueueUrl: aws.String(config.QueueURL),
+		Attributes: map[string]*string{
+			"Policy": aws.String(string(policyb)),
+		}}
+	_, err = sqsSvc.SetQueueAttributes(addPermInput)
+	if err != nil {
+		return err
+	}
+	return nil
 }
 
 func keyName(clusterName string, name string) string {
 	return clusterName + "-" + name
 }
 
-func newPolicy(queueARN string, keyname string, topicARN string) Policy {
-	return Policy{
-		Version: "2012-10-17",
-		ID:      queueARN + "/" + keyname + "-policy",
-		Statement: []Statement{
-			Statement{
-				Sid:       keyname,
-				Effect:    "Allow",
-				Principal: "*",
-				Action:    []string{"SQS:ReceiveMessage", "SQS:SendMessage"},
-				Resource:  queueARN,
-				Condition: Condition{
-					ArnEquals: ArnEquals{
-						AwsSourceArn: topicARN,
-					},
+func newPolicy(queueARN string, name string, topicARNs []string) Policy {
+	statements := []Statement{}
+	for _, topicARN := range topicARNs {
+		statements = append(statements, Statement{
+			Sid:       topicARN + " statment",
+			Effect:    "Allow",
+			Principal: "*",
+			Action:    []string{"SQS:ReceiveMessage", "SQS:SendMessage"},
+			Resource:  queueARN,
+			Condition: Condition{
+				ArnEquals: ArnEquals{
+					AwsSourceArn: topicARN,
 				},
 			},
-		},
+		})
+	}
+
+	return Policy{
+		Version:   "2012-10-17",
+		ID:        queueARN + "/" + name + "-policy",
+		Statement: statements,
 	}
 }
 
-func process(q *Queue, svc *sqs.SQS, h Handler, stopCh <-chan struct{}) error {
-	config := q.config
+func process(config *config.Config, svc *sqs.SQS, manager *queuemanager.QueueManager, ctx context.Context) error {
 	logger := config.Logger
 	for {
 		result, err := svc.ReceiveMessage(&sqs.ReceiveMessageInput{
-			QueueUrl: aws.String(q.queueURL),
+			QueueUrl: aws.String(config.QueueURL),
 			AttributeNames: aws.StringSlice([]string{
 				"SentTimestamp",
 			}),
 			MaxNumberOfMessages: aws.Int64(1),
 			MessageAttributeNames: aws.StringSlice([]string{
 				"All",
 			}),
-			WaitTimeSeconds: aws.Int64(q.timeout),
+			WaitTimeSeconds: aws.Int64(10),
 		})
 		if err != nil {
 			logger.WithError(err).Error("unable to receive messages from queue")
 			os.Exit(1)
 		}
 
 		for _, message := range result.Messages {
-			mb := &MessageBody{}
+			mb := &queuemanager.MessageBody{}
 			err := json.Unmarshal([]byte(*message.Body), mb)
 			if err != nil {
 				logger.WithError(err).Error("error parsing message body")
 			}
 			mb.ParseMessage()
 			logger.Debugf("%+v", mb)
 
+			h, ok := manager.Get(mb.TopicARN)
+			if !ok {
+				logger.Errorf("error missing handler function for topic %s", mb.TopicARN)
+			}
+
 			err = h.HandleMessage(config, mb)
 			if err != nil {
 				logger.WithError(err).Error("error processing message")
 			}
 			logger.Debugf("stackID %v updated status to %v", mb.ParsedMessage["StackId"], mb.ParsedMessage["ResourceStatus"])
 			_, err = svc.DeleteMessage(&sqs.DeleteMessageInput{
-				QueueUrl:      aws.String(q.queueURL),
+				QueueUrl:      aws.String(config.QueueURL),
 				ReceiptHandle: message.ReceiptHandle,
 			})
 
@@ -239,11 +199,5 @@ func process(q *Queue, svc *sqs.SQS, h Handler, stopCh <-chan struct{}) error {
 				os.Exit(1)
 			}
 		}
-
-		select {
-		case <-stopCh:
-			os.Exit(1)
-		default:
-		}
 	}
 }

pkg/queue/types.go

@@ -39,26 +39,3 @@ type Condition struct {
 type ArnEquals struct {
 	AwsSourceArn string `json:"aws:SourceArn"`
 }
-
-// MessageBody will parse the message from the Body of SQS
-type MessageBody struct {
-	Type               string `json:"Type"`
-	TopicArn           string `json:"TopicArn"`
-	Message            string `json:"Message"`
-	ParsedMessage      map[string]string
-	Namespace          string
-	ResourceName       string
-	ResourceProperties ResourceProperties
-	Updatable          bool
-}
-
-// ResourceProperties will wrap the ResourceProperties object
-type ResourceProperties struct {
-	Tags []Tag `json:"Tags"`
-}
-
-// Tag represents a Tag
-type Tag struct {
-	Key   string `json:"Key"`
-	Value string `json:"Value"`
-}