Uplink bucket attribution extension (#2353)

Attribution info added to the bucket info

lib/uplink/bucket_attrs_test.go

@@ -88,16 +88,18 @@ func TestPartnerBucketAttrs(t *testing.T) {
 			proj.uplinkCfg.Volatile.PartnerID = partnerID
 			got, err := proj.OpenBucket(ctx, bucketName, access)
 			require.NoError(t, err)
+			assert.Equal(t, got.bucket.Attribution, partnerID)
 
 			info, err := db.Get(ctx, consoleProject.ID, []byte(bucketName))
 			require.NoError(t, err)
 			assert.Equal(t, info.PartnerID.String(), partnerID)
 
-			// partner ID NOT set
+			// partner ID not set but bucket's attribution already set(from above)
 			proj.uplinkCfg.Volatile.PartnerID = ""
 			got, err = proj.OpenBucket(ctx, bucketName, access)
 			require.NoError(t, err)
 			defer ctx.Check(got.Close)
+			assert.Equal(t, got.bucket.Attribution, partnerID)
 		})
 }
 

lib/uplink/project.go

@@ -153,14 +153,23 @@ func (p *Project) GetBucketInfo(ctx context.Context, bucket string) (b storj.Buc
 func (p *Project) OpenBucket(ctx context.Context, bucketName string, encCtx *EncryptionCtx) (b *Bucket, err error) {
 	defer mon.Task()(&ctx)(&err)
 
-	err = p.checkBucketAttribution(ctx, bucketName)
+	bucketInfo, cfg, err := p.GetBucketInfo(ctx, bucketName)
 	if err != nil {
 		return nil, err
 	}
 
-	bucketInfo, cfg, err := p.GetBucketInfo(ctx, bucketName)
-	if err != nil {
-		return nil, err
+	// partnerID set and bucket's attribution is not set
+	if p.uplinkCfg.Volatile.PartnerID != "" && bucketInfo.Attribution == "" {
+		err = p.checkBucketAttribution(ctx, bucketName)
+		if err != nil {
+			return nil, err
+		}
+
+		// update the bucket with attribution info
+		bucketInfo, err = p.updateBucket(ctx, bucketInfo)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	encryptionScheme := cfg.EncryptionParameters.ToEncryptionScheme()
@@ -244,3 +253,17 @@ func (p *Project) checkBucketAttribution(ctx context.Context, bucketName string)
 
 	return p.metainfo.SetAttribution(ctx, bucketName, *partnerID)
 }
+
+// updateBucket updates an existing bucket's attribution info.
+func (p *Project) updateBucket(ctx context.Context, bucketInfo storj.Bucket) (bucket storj.Bucket, err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	bucket = storj.Bucket{
+		Attribution:          p.uplinkCfg.Volatile.PartnerID,
+		PathCipher:           bucketInfo.PathCipher,
+		EncryptionParameters: bucketInfo.EncryptionParameters,
+		RedundancyScheme:     bucketInfo.RedundancyScheme,
+		SegmentsSize:         bucketInfo.SegmentsSize,
+	}
+	return p.project.CreateBucket(ctx, bucketInfo.Name, &bucket)
+}

pkg/metainfo/kvmetainfo/buckets.go

@@ -19,7 +19,7 @@ import (
 	"storj.io/storj/storage"
 )
 
-// CreateBucket creates a new bucket with the specified information
+// CreateBucket creates a new bucket or updates and existing bucket with the specified information
 func (db *Project) CreateBucket(ctx context.Context, bucketName string, info *storj.Bucket) (bucketInfo storj.Bucket, err error) {
 	defer mon.Task()(&ctx)(&err)
 
@@ -67,6 +67,7 @@ func (db *Project) CreateBucket(ctx context.Context, bucketName string, info *st
 
 	r := bytes.NewReader(nil)
 	userMeta := map[string]string{
+		"attribution-to":    info.Attribution,
 		"path-enc-type":     strconv.Itoa(int(info.PathCipher)),
 		"default-seg-size":  strconv.FormatInt(info.SegmentsSize, 10),
 		"default-enc-type":  strconv.Itoa(int(info.EncryptionParameters.CipherSuite.ToCipher())),
@@ -220,6 +221,7 @@ func bucketFromMeta(ctx context.Context, bucketName string, m objects.Meta) (out
 	es := &out.EncryptionParameters
 	rs := &out.RedundancyScheme
 
+	out.Attribution = m.UserDefined["attribution-to"]
 	applySetting("path-enc-type", 16, func(v int64) { out.PathCipher = storj.Cipher(v) })
 	applySetting("default-seg-size", 64, func(v int64) { out.SegmentsSize = v })
 	applySetting("default-enc-type", 32, func(v int64) { es.CipherSuite = storj.Cipher(v).ToCipherSuite() })

pkg/storj/object.go

@@ -26,6 +26,7 @@ var (
 // Bucket contains information about a specific bucket
 type Bucket struct {
 	Name                 string
+	Attribution          string
 	Created              time.Time
 	PathCipher           Cipher
 	SegmentsSize         int64