seaweedfs/weed/notification/google_pub_sub/google_pub_sub.go

90 lines
2.4 KiB
Go
Raw Normal View History

package google_pub_sub
import (
2018-11-01 08:12:21 +00:00
"context"
"fmt"
"os"
2018-11-01 08:12:21 +00:00
"cloud.google.com/go/pubsub"
"github.com/chrislusf/seaweedfs/weed/glog"
"github.com/chrislusf/seaweedfs/weed/notification"
"github.com/chrislusf/seaweedfs/weed/util"
"github.com/golang/protobuf/proto"
"google.golang.org/api/option"
)
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
}