init pubsub

Author: conduongtong

.gitignore

@@ -1,15 +1,2 @@
-# Binaries for programs and plugins
-*.exe
-*.exe~
-*.dll
-*.so
-*.dylib
-
-# Test binary, built with `go test -c`
-*.test
-
-# Output of the go coverage tool, specifically when used with LiteIDE
-*.out
-
-# Dependency directories (remove the comment below to include it)
-# vendor/
+.vscode
+.idea

README.md

@@ -0,0 +1,15 @@
+# pubsub
+
+## Installation
+
+Please make sure to initialize a Go module before installing common-go/pubsub:
+
+```shell
+go get -u github.com/common-go/pubsub
+```
+
+Import:
+
+```go
+import "github.com/common-go/pubsub"
+```

client.go

@@ -0,0 +1,34 @@
+package pubsub
+
+import (
+	"cloud.google.com/go/pubsub"
+	"context"
+	"github.com/sirupsen/logrus"
+	"google.golang.org/api/option"
+	"os"
+)
+
+func NewPubSubClient(ctx context.Context, projectId string, keyFilename string) (*pubsub.Client, error) {
+	if len(keyFilename) > 0 && existFile(keyFilename) {
+		if logrus.IsLevelEnabled(logrus.InfoLevel) {
+			logrus.Info("key file exists")
+		}
+		return pubsub.NewClient(ctx, projectId, option.WithCredentialsFile(keyFilename))
+	} else {
+		if logrus.IsLevelEnabled(logrus.WarnLevel) && len(keyFilename) > 0{
+			logrus.Warn("key file doesn't exists")
+		}
+		return pubsub.NewClient(ctx, projectId)
+	}
+}
+
+func existFile(filename string) bool {
+	if _, err := os.Stat(filename); err == nil {
+		return true
+	} else if os.IsNotExist(err) {
+		return false
+	} else {
+		logrus.Error(err)
+	}
+	return false
+}

client_config.go

@@ -0,0 +1,6 @@
+package pubsub
+
+type ClientConfig struct {
+	ProjectId   string `mapstructure:"project_id"`
+	KeyFilename string `mapstructure:"key_filename"`
+}

consumer.go

@@ -0,0 +1,58 @@
+package pubsub
+
+import (
+	"cloud.google.com/go/pubsub"
+	"context"
+	"github.com/common-go/mq"
+	"github.com/sirupsen/logrus"
+)
+
+type Consumer struct {
+	Client       *pubsub.Client
+	Subscription *pubsub.Subscription
+	AckOnConsume bool
+}
+
+func NewConsumer(client *pubsub.Client, subscriptionId string, c SubscriptionConfig, ackOnConsume bool) *Consumer {
+	subscription := client.Subscription(subscriptionId)
+	return &Consumer{Client: client, Subscription: ConfigureSubscription(subscription, c), AckOnConsume: ackOnConsume}
+}
+
+func NewConsumerByConfig(ctx context.Context, c ConsumerConfig, ackOnConsume bool) (*Consumer, error) {
+	client, err := NewPubSubClient(ctx, c.Client.ProjectId, c.Client.KeyFilename)
+	if err != nil {
+		return nil, err
+	}
+	return NewConsumer(client, c.SubscriptionId, c.SubscriptionConfig, ackOnConsume), nil
+}
+
+func ConfigureSubscription(subscription *pubsub.Subscription, c SubscriptionConfig) *pubsub.Subscription {
+	if c.MaxOutstandingMessages > 0 {
+		subscription.ReceiveSettings.MaxOutstandingMessages = c.MaxOutstandingMessages
+	}
+	if c.NumGoroutines > 0 {
+		subscription.ReceiveSettings.NumGoroutines = c.NumGoroutines
+	}
+	return subscription
+}
+
+func (c *Consumer) Consume(ctx context.Context, caller mq.ConsumerCaller) {
+	er1 := c.Subscription.Receive(ctx, func(ctx2 context.Context, m *pubsub.Message) {
+		if logrus.IsLevelEnabled(logrus.DebugLevel) {
+			logrus.Debugf("Received message: %s", m.Data)
+		}
+		message := mq.Message{
+			Id:         m.ID,
+			Data:       m.Data,
+			Attributes: m.Attributes,
+			Raw:        m,
+		}
+		if c.AckOnConsume {
+			m.Ack()
+		}
+		caller.Call(ctx2, &message, nil)
+	})
+	if er1 != nil {
+		caller.Call(ctx, nil, er1)
+	}
+}

consumer_config.go

@@ -0,0 +1,12 @@
+package pubsub
+
+type ConsumerConfig struct {
+	SubscriptionId     string             `mapstructure:"subscription_id"`
+	Client             ClientConfig       `mapstructure:"client"`
+	SubscriptionConfig SubscriptionConfig `mapstructure:"subscription"`
+}
+
+type SubscriptionConfig struct {
+	MaxOutstandingMessages int `mapstructure:"max_outstanding_messages"`
+	NumGoroutines          int `mapstructure:"num_goroutines"`
+}

producer.go

@@ -0,0 +1,75 @@
+package pubsub
+
+import (
+	"context"
+	"time"
+
+	"cloud.google.com/go/iam"
+	"cloud.google.com/go/pubsub"
+	"github.com/sirupsen/logrus"
+)
+
+var CheckTopicPermission = CheckPermission
+
+type Producer struct {
+	Client *pubsub.Client
+	Topic  *pubsub.Topic
+}
+
+func NewProducer(ctx context.Context, client *pubsub.Client, topicId string, c TopicConfig) *Producer {
+	topic := client.Topic(topicId)
+	CheckTopicPermission(ctx, topic.IAM(), "pubsub.topics.publish")
+	return &Producer{Client: client, Topic: ConfigureTopic(topic, c)}
+}
+
+func NewProducerByConfig(ctx context.Context, c ProducerConfig) (*Producer, error) {
+	client, err := NewPubSubClient(ctx, c.Client.ProjectId, c.Client.KeyFilename)
+	if err != nil {
+		return nil, err
+	}
+	return NewProducer(ctx, client, c.TopicId, c.Topic), nil
+}
+
+func ConfigureTopic(topic *pubsub.Topic, c TopicConfig) *pubsub.Topic {
+	if c.CountThreshold > 0 {
+		topic.PublishSettings.DelayThreshold = time.Duration(c.CountThreshold) * time.Millisecond
+	}
+	if c.DelayThreshold > 0 {
+		topic.PublishSettings.CountThreshold = c.DelayThreshold
+	}
+	if c.ByteThreshold > 0 {
+		topic.PublishSettings.ByteThreshold = c.ByteThreshold
+	}
+	if c.NumGoroutines > 0 {
+		topic.PublishSettings.NumGoroutines = c.NumGoroutines
+	}
+	return topic
+}
+
+func (c *Producer) Produce(ctx context.Context, data []byte, messageAttributes *map[string]string) (string, error) {
+	msg := &pubsub.Message{
+		Data: data,
+	}
+
+	if messageAttributes != nil {
+		msg.Attributes = *messageAttributes
+	}
+
+	publishResult := c.Topic.Publish(ctx, msg)
+	return publishResult.Get(ctx)
+}
+
+func CheckPermission(ctx0 context.Context, iam *iam.Handle, permission string) {
+	ctx, _ := context.WithTimeout(ctx0, 30*time.Second)
+
+	if logrus.IsLevelEnabled(logrus.InfoLevel) {
+		logrus.Infof("Checking permission: %s", permission)
+	}
+	if permissions, err := iam.TestPermissions(ctx, []string{permission}); err != nil {
+		logrus.Fatalf("Can't check permission %s: %v", permission, err)
+	} else if len(permissions) > 0 && permissions[0] == permission {
+		logrus.Warn("Permission %s valid", permission)
+	} else {
+		logrus.Fatalf("Permission %s invalid", permission)
+	}
+}

producer_config.go

@@ -0,0 +1,14 @@
+package pubsub
+
+type ProducerConfig struct {
+	TopicId string       `mapstructure:"topic_id"`
+	Client  ClientConfig `mapstructure:"client"`
+	Topic   TopicConfig  `mapstructure:"topic"`
+}
+
+type TopicConfig struct {
+	DelayThreshold int `mapstructure:"delay_threshold"` // MaxMessages
+	CountThreshold int `mapstructure:"count_threshold"` // MaxMilliseconds
+	ByteThreshold  int `mapstructure:"byte_threshold"`  // MaxBytes
+	NumGoroutines  int `mapstructure:"num_goroutines"`
+}

pubsub_health_service.go

@@ -0,0 +1,63 @@
+package pubsub
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"cloud.google.com/go/pubsub"
+	"github.com/sirupsen/logrus"
+)
+
+type PermissionType int
+
+const (
+	PermissionPublish   PermissionType = 0
+	PermissionSubscribe PermissionType = 1
+)
+
+type PubSubHealthService struct {
+	name           string
+	client         *pubsub.Client
+	timeout        time.Duration
+	permissionType PermissionType
+	resourceId     string
+}
+
+func NewPubSubHealthService(name string, client *pubsub.Client, timeout time.Duration, permissionType PermissionType, resourceId string) *PubSubHealthService {
+	return &PubSubHealthService{name, client, timeout, permissionType, resourceId}
+}
+
+func (h *PubSubHealthService) Name() string {
+	return h.name
+}
+
+func (h *PubSubHealthService) Check(ctx context.Context) (map[string]interface{}, error) {
+	res := make(map[string]interface{})
+	var permissions []string
+	var err error
+
+	timeoutCtx, _ := context.WithTimeout(ctx, h.timeout)
+	if h.permissionType == PermissionPublish {
+		permissions, err = h.client.Topic(h.resourceId).IAM().TestPermissions(timeoutCtx, []string{"pubsub.topics.publish"})
+	} else if h.permissionType == PermissionSubscribe {
+		permissions, err = h.client.Subscription(h.resourceId).IAM().TestPermissions(timeoutCtx, []string{"pubsub.subscriptions.consume"})
+	}
+
+	if err != nil {
+		logrus.Errorf("Can't TestPermissions %h: %h", h.resourceId, err.Error())
+		return res, err
+	} else if len(permissions) != 1 {
+		return res, fmt.Errorf("invalid permissions: %v", permissions)
+	} else {
+		return res, nil
+	}
+}
+
+func (h *PubSubHealthService) Build(ctx context.Context, data map[string]interface{}, err error) map[string]interface{} {
+	if err == nil {
+		return data
+	}
+	data["error"] = err.Error()
+	return data
+}