gcs_sink.go 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. package gcssink
  2. import (
  3. "context"
  4. "fmt"
  5. "github.com/chrislusf/seaweedfs/weed/replication/repl_util"
  6. "os"
  7. "cloud.google.com/go/storage"
  8. "google.golang.org/api/option"
  9. "github.com/chrislusf/seaweedfs/weed/filer"
  10. "github.com/chrislusf/seaweedfs/weed/glog"
  11. "github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
  12. "github.com/chrislusf/seaweedfs/weed/replication/sink"
  13. "github.com/chrislusf/seaweedfs/weed/replication/source"
  14. "github.com/chrislusf/seaweedfs/weed/util"
  15. )
  16. type GcsSink struct {
  17. client *storage.Client
  18. bucket string
  19. dir string
  20. filerSource *source.FilerSource
  21. isIncremental bool
  22. }
  23. func init() {
  24. sink.Sinks = append(sink.Sinks, &GcsSink{})
  25. }
  26. func (g *GcsSink) GetName() string {
  27. return "google_cloud_storage"
  28. }
  29. func (g *GcsSink) GetSinkToDirectory() string {
  30. return g.dir
  31. }
  32. func (g *GcsSink) IsIncremental() bool {
  33. return g.isIncremental
  34. }
  35. func (g *GcsSink) Initialize(configuration util.Configuration, prefix string) error {
  36. g.isIncremental = configuration.GetBool(prefix + "is_incremental")
  37. return g.initialize(
  38. configuration.GetString(prefix+"google_application_credentials"),
  39. configuration.GetString(prefix+"bucket"),
  40. configuration.GetString(prefix+"directory"),
  41. )
  42. }
  43. func (g *GcsSink) SetSourceFiler(s *source.FilerSource) {
  44. g.filerSource = s
  45. }
  46. func (g *GcsSink) initialize(google_application_credentials, bucketName, dir string) error {
  47. g.bucket = bucketName
  48. g.dir = dir
  49. // Creates a client.
  50. if google_application_credentials == "" {
  51. var found bool
  52. google_application_credentials, found = os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
  53. if !found {
  54. glog.Fatalf("need to specific GOOGLE_APPLICATION_CREDENTIALS env variable or google_application_credentials in replication.toml")
  55. }
  56. }
  57. client, err := storage.NewClient(context.Background(), option.WithCredentialsFile(google_application_credentials))
  58. if err != nil {
  59. glog.Fatalf("Failed to create client: %v", err)
  60. }
  61. g.client = client
  62. return nil
  63. }
  64. func (g *GcsSink) DeleteEntry(key string, isDirectory, deleteIncludeChunks bool, signatures []int32) error {
  65. if isDirectory {
  66. key = key + "/"
  67. }
  68. if err := g.client.Bucket(g.bucket).Object(key).Delete(context.Background()); err != nil {
  69. return fmt.Errorf("gcs delete %s%s: %v", g.bucket, key, err)
  70. }
  71. return nil
  72. }
  73. func (g *GcsSink) CreateEntry(key string, entry *filer_pb.Entry, signatures []int32) error {
  74. if entry.IsDirectory {
  75. return nil
  76. }
  77. totalSize := filer.FileSize(entry)
  78. chunkViews := filer.ViewFromChunks(g.filerSource.LookupFileId, entry.Chunks, 0, int64(totalSize))
  79. wc := g.client.Bucket(g.bucket).Object(key).NewWriter(context.Background())
  80. defer wc.Close()
  81. writeFunc := func(data []byte) error {
  82. _, writeErr := wc.Write(data)
  83. return writeErr
  84. }
  85. if err := repl_util.CopyFromChunkViews(chunkViews, g.filerSource, writeFunc); err != nil {
  86. return err
  87. }
  88. return nil
  89. }
  90. func (g *GcsSink) UpdateEntry(key string, oldEntry *filer_pb.Entry, newParentPath string, newEntry *filer_pb.Entry, deleteIncludeChunks bool, signatures []int32) (foundExistingEntry bool, err error) {
  91. return true, g.CreateEntry(key, newEntry, signatures)
  92. }