You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

115 lines
3.1 KiB

6 years ago
  1. package sub
  2. import (
  3. "context"
  4. "fmt"
  5. "os"
  6. "cloud.google.com/go/pubsub"
  7. "github.com/golang/protobuf/proto"
  8. "github.com/seaweedfs/seaweedfs/weed/glog"
  9. "github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
  10. "github.com/seaweedfs/seaweedfs/weed/util"
  11. "google.golang.org/api/option"
  12. )
  13. func init() {
  14. NotificationInputs = append(NotificationInputs, &GooglePubSubInput{})
  15. }
  16. type GooglePubSubInput struct {
  17. sub *pubsub.Subscription
  18. topicName string
  19. messageChan chan *pubsub.Message
  20. }
  21. func (k *GooglePubSubInput) GetName() string {
  22. return "google_pub_sub"
  23. }
  24. func (k *GooglePubSubInput) Initialize(configuration util.Configuration, prefix string) error {
  25. glog.V(0).Infof("notification.google_pub_sub.project_id: %v", configuration.GetString(prefix+"project_id"))
  26. glog.V(0).Infof("notification.google_pub_sub.topic: %v", configuration.GetString(prefix+"topic"))
  27. return k.initialize(
  28. configuration.GetString(prefix+"google_application_credentials"),
  29. configuration.GetString(prefix+"project_id"),
  30. configuration.GetString(prefix+"topic"),
  31. )
  32. }
  33. func (k *GooglePubSubInput) initialize(google_application_credentials, projectId, topicName string) (err error) {
  34. ctx := context.Background()
  35. // Creates a client.
  36. if google_application_credentials == "" {
  37. var found bool
  38. google_application_credentials, found = os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
  39. if !found {
  40. glog.Fatalf("need to specific GOOGLE_APPLICATION_CREDENTIALS env variable or google_application_credentials in filer.toml")
  41. }
  42. }
  43. client, err := pubsub.NewClient(ctx, projectId, option.WithCredentialsFile(google_application_credentials))
  44. if err != nil {
  45. glog.Fatalf("Failed to create client: %v", err)
  46. }
  47. k.topicName = topicName
  48. topic := client.Topic(topicName)
  49. if exists, err := topic.Exists(ctx); err == nil {
  50. if !exists {
  51. topic, err = client.CreateTopic(ctx, topicName)
  52. if err != nil {
  53. glog.Fatalf("Failed to create topic %s: %v", topicName, err)
  54. }
  55. }
  56. } else {
  57. glog.Fatalf("Failed to check topic %s: %v", topicName, err)
  58. }
  59. subscriptionName := "seaweedfs_sub"
  60. k.sub = client.Subscription(subscriptionName)
  61. if exists, err := k.sub.Exists(ctx); err == nil {
  62. if !exists {
  63. k.sub, err = client.CreateSubscription(ctx, subscriptionName, pubsub.SubscriptionConfig{Topic: topic})
  64. if err != nil {
  65. glog.Fatalf("Failed to create subscription %s: %v", subscriptionName, err)
  66. }
  67. }
  68. } else {
  69. glog.Fatalf("Failed to check subscription %s: %v", topicName, err)
  70. }
  71. k.messageChan = make(chan *pubsub.Message, 1)
  72. go k.sub.Receive(ctx, func(ctx context.Context, m *pubsub.Message) {
  73. k.messageChan <- m
  74. })
  75. return err
  76. }
  77. func (k *GooglePubSubInput) ReceiveMessage() (key string, message *filer_pb.EventNotification, onSuccessFn func(), onFailureFn func(), err error) {
  78. m := <-k.messageChan
  79. onSuccessFn = func() {
  80. m.Ack()
  81. }
  82. onFailureFn = func() {
  83. m.Nack()
  84. }
  85. // process the message
  86. key = m.Attributes["key"]
  87. message = &filer_pb.EventNotification{}
  88. err = proto.Unmarshal(m.Data, message)
  89. if err != nil {
  90. err = fmt.Errorf("unmarshal message from google pubsub %s: %v", k.topicName, err)
  91. return
  92. }
  93. return
  94. }