satellite/metabase: add GetProjectSegmentCount method

This new method will be used with mechanism to limit number
of segments per project, similar to limiting buckets or bandwidth.
This is only one of multiple changes we will do to implement this
limitation.

Change-Id: Ia516c2a006ad1d7b4431d780679be9d809848f4b

satellite/metabase/metabasetest/test.go

@@ -711,3 +711,18 @@ func (step FinishMoveObject) Check(ctx *testcontext.Context, t testing.TB, db *m
 	err := db.FinishMoveObject(ctx, step.Opts)
 	checkError(t, err, step.ErrClass, step.ErrText)
 }
+
+// GetProjectSegmentCount is for testing metabase.GetProjectSegmentCount.
+type GetProjectSegmentCount struct {
+	Opts     metabase.GetProjectSegmentCount
+	Result   int64
+	ErrClass *errs.Class
+	ErrText  string
+}
+
+// Check runs the test.
+func (step GetProjectSegmentCount) Check(ctx *testcontext.Context, t testing.TB, db *metabase.DB) {
+	result, err := db.GetProjectSegmentCount(ctx, step.Opts)
+	checkError(t, err, step.ErrClass, step.ErrText)
+	require.Equal(t, step.Result, result)
+}

satellite/metabase/project.go

@@ -0,0 +1,56 @@
+// Copyright (C) 2021 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package metabase
+
+import (
+	"context"
+	"time"
+
+	"storj.io/common/uuid"
+)
+
+// GetProjectSegmentCount contains arguments necessary for fetching an information
+// about project segment count.
+type GetProjectSegmentCount struct {
+	ProjectID uuid.UUID
+
+	AsOfSystemTime     time.Time
+	AsOfSystemInterval time.Duration
+}
+
+// Verify verifies reqest fields.
+func (g *GetProjectSegmentCount) Verify() error {
+	if g.ProjectID.IsZero() {
+		return ErrInvalidRequest.New("ProjectID missing")
+	}
+	return nil
+}
+
+// GetProjectSegmentCount returns number of segments that specified project has.
+func (db *DB) GetProjectSegmentCount(ctx context.Context, opts GetProjectSegmentCount) (_ int64, err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	if err := opts.Verify(); err != nil {
+		return 0, err
+	}
+
+	var segmentsCount *int64
+	err = db.db.QueryRowContext(ctx, `
+		SELECT
+			sum(segment_count)
+		FROM objects
+		`+db.asOfTime(opts.AsOfSystemTime, opts.AsOfSystemInterval)+`
+		WHERE
+			project_id = $1
+		`, opts.ProjectID).Scan(&segmentsCount)
+	if err != nil {
+		return 0, Error.New("unable to query project segment count: %w", err)
+	}
+
+	if segmentsCount == nil {
+		return 0, Error.New("project not found: %s", opts.ProjectID)
+	}
+
+	return *segmentsCount, nil
+}

satellite/metabase/project_test.go

@@ -0,0 +1,133 @@
+// Copyright (C) 2021 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package metabase_test
+
+import (
+	"testing"
+	"time"
+
+	"storj.io/common/testcontext"
+	"storj.io/common/testrand"
+	"storj.io/common/uuid"
+	"storj.io/storj/satellite/metabase"
+	"storj.io/storj/satellite/metabase/metabasetest"
+)
+
+func TestGetProjectSegmentsCount(t *testing.T) {
+	metabasetest.Run(t, func(ctx *testcontext.Context, t *testing.T, db *metabase.DB) {
+		obj := metabasetest.RandObjectStream()
+
+		t.Run("ProjectID missing", func(t *testing.T) {
+			defer metabasetest.DeleteAll{}.Check(ctx, t, db)
+
+			metabasetest.GetProjectSegmentCount{
+				Opts:     metabase.GetProjectSegmentCount{},
+				ErrClass: &metabase.ErrInvalidRequest,
+				ErrText:  "ProjectID missing",
+			}.Check(ctx, t, db)
+
+			metabasetest.Verify{}.Check(ctx, t, db)
+		})
+
+		t.Run("empty database", func(t *testing.T) {
+			defer metabasetest.DeleteAll{}.Check(ctx, t, db)
+
+			metabasetest.GetProjectSegmentCount{
+				Opts: metabase.GetProjectSegmentCount{
+					ProjectID: obj.ProjectID,
+				},
+				ErrClass: &metabase.Error,
+				ErrText:  "project not found: " + obj.ProjectID.String(),
+			}.Check(ctx, t, db)
+		})
+
+		t.Run("object without segments", func(t *testing.T) {
+			defer metabasetest.DeleteAll{}.Check(ctx, t, db)
+
+			metabasetest.CreateObject(ctx, t, db, obj, 0)
+
+			metabasetest.GetProjectSegmentCount{
+				Opts: metabase.GetProjectSegmentCount{
+					ProjectID: obj.ProjectID,
+				},
+				Result: 0,
+			}.Check(ctx, t, db)
+		})
+
+		t.Run("object with segments", func(t *testing.T) {
+			defer metabasetest.DeleteAll{}.Check(ctx, t, db)
+
+			metabasetest.CreateObject(ctx, t, db, obj, 1)
+
+			metabasetest.GetProjectSegmentCount{
+				Opts: metabase.GetProjectSegmentCount{
+					ProjectID: obj.ProjectID,
+				},
+				Result: 1,
+			}.Check(ctx, t, db)
+		})
+
+		t.Run("object with segments (as of system time)", func(t *testing.T) {
+			defer metabasetest.DeleteAll{}.Check(ctx, t, db)
+
+			metabasetest.CreateObject(ctx, t, db, obj, 1)
+
+			metabasetest.GetProjectSegmentCount{
+				Opts: metabase.GetProjectSegmentCount{
+					ProjectID:          obj.ProjectID,
+					AsOfSystemTime:     time.Now(),
+					AsOfSystemInterval: time.Millisecond,
+				},
+				Result: 1,
+			}.Check(ctx, t, db)
+		})
+
+		t.Run("multiple projects", func(t *testing.T) {
+			defer metabasetest.DeleteAll{}.Check(ctx, t, db)
+
+			project1 := testrand.UUID()
+			project2 := testrand.UUID()
+			project3 := testrand.UUID()
+
+			type Object struct {
+				ProjectID uuid.UUID
+				Segments  byte
+			}
+
+			objects := []Object{
+				{project1, 1},
+				{project1, 4},
+
+				{project2, 5},
+				{project2, 3},
+				{project2, 1},
+
+				{project3, 2},
+				{project3, 0},
+			}
+
+			for _, object := range objects {
+				obj := metabasetest.RandObjectStream()
+				obj.ProjectID = object.ProjectID
+				metabasetest.CreateObject(ctx, t, db, obj, object.Segments)
+			}
+
+			// expected number of segments per project
+			projects := map[uuid.UUID]int{
+				project1: 5,
+				project2: 9,
+				project3: 2,
+			}
+
+			for project, segments := range projects {
+				metabasetest.GetProjectSegmentCount{
+					Opts: metabase.GetProjectSegmentCount{
+						ProjectID: project,
+					},
+					Result: int64(segments),
+				}.Check(ctx, t, db)
+			}
+		})
+	})
+}