upstream-seaweedfs/weed/notification/google_pub_sub/google_pub_sub.go

package google_pub_sub

import (
	"context"
	"fmt"
	"os"

	"cloud.google.com/go/pubsub"
	"github.com/seaweedfs/seaweedfs/weed/glog"
	"github.com/seaweedfs/seaweedfs/weed/notification"
	"github.com/seaweedfs/seaweedfs/weed/util"
	"google.golang.org/api/option"
	"google.golang.org/protobuf/proto"
)

func init() {
	notification.MessageQueues = append(notification.MessageQueues, &GooglePubSub{})
}

type GooglePubSub struct {
	topic *pubsub.Topic
}

func (k *GooglePubSub) GetName() string {
	return "google_pub_sub"
}

func (k *GooglePubSub) Initialize(configuration util.Configuration, prefix string) (err error) {
	glog.V(0).Infof("notification.google_pub_sub.project_id: %v", configuration.GetString(prefix+"project_id"))
	glog.V(0).Infof("notification.google_pub_sub.topic: %v", configuration.GetString(prefix+"topic"))
	return k.initialize(
		configuration.GetString(prefix+"google_application_credentials"),
		configuration.GetString(prefix+"project_id"),
		configuration.GetString(prefix+"topic"),
	)
}

func (k *GooglePubSub) initialize(google_application_credentials, projectId, topicName string) (err error) {

	ctx := context.Background()
	// Creates a client.
	if google_application_credentials == "" {
		var found bool
		google_application_credentials, found = os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
		if !found {
			glog.Fatalf("need to specific GOOGLE_APPLICATION_CREDENTIALS env variable or google_application_credentials in filer.toml")
		}
	}

	client, err := pubsub.NewClient(ctx, projectId, option.WithCredentialsFile(google_application_credentials))
	if err != nil {
		glog.Fatalf("Failed to create client: %v", err)
	}

	k.topic = client.Topic(topicName)
	if exists, err := k.topic.Exists(ctx); err == nil {
		if !exists {
			k.topic, err = client.CreateTopic(ctx, topicName)
			if err != nil {
				glog.Fatalf("Failed to create topic %s: %v", topicName, err)
			}
		}
	} else {
		glog.Fatalf("Failed to check topic %s: %v", topicName, err)
	}

	return nil
}

func (k *GooglePubSub) SendMessage(key string, message proto.Message) (err error) {

	bytes, err := proto.Marshal(message)
	if err != nil {
		return
	}

	ctx := context.Background()
	result := k.topic.Publish(ctx, &pubsub.Message{
		Data:       bytes,
		Attributes: map[string]string{"key": key},
	})

	_, err = result.Get(ctx)
	if err != nil {
		return fmt.Errorf("send message to google pub sub %s: %v", k.topic.String(), err)
	}

	return nil
}