satellite/satellitedb: make GetProjectTotalByPartner always return pa…

…rtnered usage

There was a problem that GetProjectTotalByPartner would return non-partnered empty usage for a project where buckets haven't been created yet or all the buckets have been deleted.

Issue:
storj/storj-private#581

Change-Id: Ia8837f8a69485d409a64afcbf68d4afbbe88f132

satellite/payments/stripe/service_test.go

@@ -27,6 +27,7 @@ import (
 	"storj.io/storj/private/testplanet"
 	"storj.io/storj/satellite"
 	"storj.io/storj/satellite/accounting"
+	"storj.io/storj/satellite/attribution"
 	"storj.io/storj/satellite/buckets"
 	"storj.io/storj/satellite/console"
 	"storj.io/storj/satellite/metabase"
@@ -1473,6 +1474,13 @@ func TestProjectUsagePrice(t *testing.T) {
 				})
 				require.NoError(t, err)
 
+				_, err = sat.DB.Attribution().Insert(ctx, &attribution.Info{
+					ProjectID:  project.ID,
+					BucketName: []byte(bucket.Name),
+					UserAgent:  tt.userAgent,
+				})
+				require.NoError(t, err)
+
 				err = sat.DB.Orders().UpdateBucketBandwidthSettle(ctx, project.ID, []byte(bucket.Name),
 					pb.PieceAction_GET, memory.TB.Int64(), 0, period)
 				require.NoError(t, err)

satellite/satellitedb/dbx/project.dbx

@@ -108,6 +108,10 @@ read one (
     select project.default_versioning
     where project.id = ?
 )
+read one (
+    select project.user_agent
+    where project.id = ?
+)
 
 read all (
     select project

satellite/satellitedb/dbx/satellitedb.dbx.go

@@ -16257,6 +16257,28 @@ func (obj *pgxImpl) Get_Project_DefaultVersioning_By_Id(ctx context.Context,
 
 }
 
+func (obj *pgxImpl) Get_Project_UserAgent_By_Id(ctx context.Context,
+	project_id Project_Id_Field) (
+	row *UserAgent_Row, err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	var __embed_stmt = __sqlbundle_Literal("SELECT projects.user_agent FROM projects WHERE projects.id = ?")
+
+	var __values []interface{}
+	__values = append(__values, project_id.value())
+
+	var __stmt = __sqlbundle_Render(obj.dialect, __embed_stmt)
+	obj.logStmt(__stmt, __values...)
+
+	row = &UserAgent_Row{}
+	err = obj.queryRowContext(ctx, __stmt, __values...).Scan(&row.UserAgent)
+	if err != nil {
+		return (*UserAgent_Row)(nil), obj.makeErr(err)
+	}
+	return row, nil
+
+}
+
 func (obj *pgxImpl) All_Project(ctx context.Context) (
 	rows []*Project, err error) {
 	defer mon.Task()(&ctx)(&err)
@@ -24835,6 +24857,28 @@ func (obj *pgxcockroachImpl) Get_Project_DefaultVersioning_By_Id(ctx context.Con
 
 }
 
+func (obj *pgxcockroachImpl) Get_Project_UserAgent_By_Id(ctx context.Context,
+	project_id Project_Id_Field) (
+	row *UserAgent_Row, err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	var __embed_stmt = __sqlbundle_Literal("SELECT projects.user_agent FROM projects WHERE projects.id = ?")
+
+	var __values []interface{}
+	__values = append(__values, project_id.value())
+
+	var __stmt = __sqlbundle_Render(obj.dialect, __embed_stmt)
+	obj.logStmt(__stmt, __values...)
+
+	row = &UserAgent_Row{}
+	err = obj.queryRowContext(ctx, __stmt, __values...).Scan(&row.UserAgent)
+	if err != nil {
+		return (*UserAgent_Row)(nil), obj.makeErr(err)
+	}
+	return row, nil
+
+}
+
 func (obj *pgxcockroachImpl) All_Project(ctx context.Context) (
 	rows []*Project, err error) {
 	defer mon.Task()(&ctx)(&err)
@@ -30170,6 +30214,10 @@ type Methods interface {
 		project_id Project_Id_Field) (
 		row *UsageLimit_Row, err error)
 
+	Get_Project_UserAgent_By_Id(ctx context.Context,
+		project_id Project_Id_Field) (
+		row *UserAgent_Row, err error)
+
 	Get_Project_UserSpecifiedBandwidthLimit_By_Id(ctx context.Context,
 		project_id Project_Id_Field) (
 		row *UserSpecifiedBandwidthLimit_Row, err error)

satellite/satellitedb/projectaccounting.go

@@ -593,28 +593,19 @@ func (db *ProjectAccounting) GetProjectTotalByPartner(ctx context.Context, proje
 	usages = make(map[string]accounting.ProjectUsage)
 
 	for _, bucket := range bucketNames {
-		userAgentRow, err := db.db.Get_BucketMetainfo_UserAgent_By_ProjectId_And_Name(ctx,
-			dbx.BucketMetainfo_ProjectId(projectID[:]),
-			dbx.BucketMetainfo_Name([]byte(bucket)))
+		valueAttr, err := db.db.Get_ValueAttribution_By_ProjectId_And_BucketName(ctx,
+			dbx.ValueAttribution_ProjectId(projectID[:]),
+			dbx.ValueAttribution_BucketName([]byte(bucket)))
 		if err != nil && !errors.Is(err, sql.ErrNoRows) {
 			return nil, err
 		}
 
 		var partner string
-		if userAgentRow != nil && userAgentRow.UserAgent != nil {
-			entries, err := useragent.ParseEntries(userAgentRow.UserAgent)
+		if valueAttr != nil && valueAttr.UserAgent != nil {
+			partner, err = tryFindPartnerByUserAgent(valueAttr.UserAgent, partnerNames)
 			if err != nil {
 				return nil, err
 			}
-
-			if len(entries) != 0 {
-				for _, iterPartner := range partnerNames {
-					if entries[0].Product == iterPartner {
-						partner = iterPartner
-						break
-					}
-				}
-			}
 		}
 		if _, ok := usages[partner]; !ok {
 			usages[partner] = accounting.ProjectUsage{Since: since, Before: before}
@@ -671,9 +662,43 @@ func (db *ProjectAccounting) GetProjectTotalByPartner(ctx context.Context, proje
 		usages[partner] = usage
 	}
 
+	// We search for project user_agent for cases when buckets haven't been created yet.
+	if len(usages) == 0 {
+		userAgentRow, err := db.db.Get_Project_UserAgent_By_Id(ctx, dbx.Project_Id(projectID[:]))
+		if err != nil && !errors.Is(err, sql.ErrNoRows) {
+			return nil, err
+		}
+		if userAgentRow != nil && userAgentRow.UserAgent != nil {
+			partner, err := tryFindPartnerByUserAgent(userAgentRow.UserAgent, partnerNames)
+			if err != nil {
+				return nil, err
+			}
+			usages[partner] = accounting.ProjectUsage{}
+		}
+	}
+
 	return usages, nil
 }
 
+func tryFindPartnerByUserAgent(userAgent []byte, partnerNames []string) (string, error) {
+	entries, err := useragent.ParseEntries(userAgent)
+	if err != nil {
+		return "", err
+	}
+
+	var partner string
+	if len(entries) != 0 {
+		for _, iterPartner := range partnerNames {
+			if entries[0].Product == iterPartner {
+				partner = iterPartner
+				break
+			}
+		}
+	}
+
+	return partner, nil
+}
+
 // GetBucketUsageRollups retrieves summed usage rollups for every bucket of particular project for a given period.
 func (db *ProjectAccounting) GetBucketUsageRollups(ctx context.Context, projectID uuid.UUID, since, before time.Time) (_ []accounting.BucketUsageRollup, err error) {
 	defer mon.Task()(&ctx)(&err)

satellite/satellitedb/projectaccounting_test.go

@@ -17,6 +17,7 @@ import (
 	"storj.io/storj/private/testplanet"
 	"storj.io/storj/satellite"
 	"storj.io/storj/satellite/accounting"
+	"storj.io/storj/satellite/attribution"
 	"storj.io/storj/satellite/buckets"
 	"storj.io/storj/satellite/console"
 	"storj.io/storj/satellite/metabase"
@@ -316,6 +317,13 @@ func TestGetProjectTotalByPartner(t *testing.T) {
 				_, err := sat.DB.Buckets().CreateBucket(ctx, bucket)
 				require.NoError(t, err)
 
+				_, err = sat.DB.Attribution().Insert(ctx, &attribution.Info{
+					ProjectID:  project.ID,
+					BucketName: []byte(bucket.Name),
+					UserAgent:  bucket.UserAgent,
+				})
+				require.NoError(t, err)
+
 				// We use multiple tallies and rollups to ensure that
 				// GetProjectTotalByPartner is capable of summing them.
 				for i := 0; i <= tallyRollupCount; i++ {