satellite/metainfo/metabase: add ListStreamPositions

Change-Id: I0b5a5803be8987e7cafbcc80460217d9eb8f0238

satellite/metainfo/metabase/list_segments.go

@@ -93,3 +93,73 @@ func (db *DB) ListSegments(ctx context.Context, opts ListSegments) (result ListS
 
 	return result, nil
 }
+
+// ListStreamPositions contains arguments necessary for listing stream segments.
+type ListStreamPositions struct {
+	StreamID uuid.UUID
+	Cursor   SegmentPosition
+	Limit    int
+}
+
+// ListStreamPositionsResult result of listing segments.
+type ListStreamPositionsResult struct {
+	Segments []SegmentPositionInfo
+	More     bool
+}
+
+// SegmentPositionInfo contains information for segment position.
+type SegmentPositionInfo struct {
+	Position  SegmentPosition
+	PlainSize int32
+}
+
+// ListStreamPositions lists specified stream segment positions.
+func (db *DB) ListStreamPositions(ctx context.Context, opts ListStreamPositions) (result ListStreamPositionsResult, err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	if opts.StreamID.IsZero() {
+		return ListStreamPositionsResult{}, ErrInvalidRequest.New("StreamID missing")
+	}
+
+	if opts.Limit < 0 {
+		return ListStreamPositionsResult{}, ErrInvalidRequest.New("Invalid limit: %d", opts.Limit)
+	}
+
+	if opts.Limit == 0 || opts.Limit > MaxListLimit {
+		opts.Limit = MaxListLimit
+	}
+
+	err = withRows(db.db.Query(ctx, `
+		SELECT
+			position, plain_size
+		FROM segments
+		WHERE
+			stream_id = $1 AND
+			($2 = 0::INT8 OR position > $2)
+		ORDER BY position ASC
+		LIMIT $3
+	`, opts.StreamID, opts.Cursor, opts.Limit+1))(func(rows tagsql.Rows) error {
+		for rows.Next() {
+			var segment SegmentPositionInfo
+			err = rows.Scan(&segment.Position, &segment.PlainSize)
+			if err != nil {
+				return Error.New("failed to scan segments: %w", err)
+			}
+			result.Segments = append(result.Segments, segment)
+		}
+		return nil
+	})
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return ListStreamPositionsResult{}, nil
+		}
+		return ListStreamPositionsResult{}, Error.New("unable to fetch object segments: %w", err)
+	}
+
+	if len(result.Segments) > opts.Limit {
+		result.More = true
+		result.Segments = result.Segments[:len(result.Segments)-1]
+	}
+
+	return result, nil
+}

satellite/metainfo/metabase/list_segments_test.go

@@ -259,3 +259,257 @@ func TestListSegments(t *testing.T) {
 		})
 	})
 }
+
+func TestListStreamPositions(t *testing.T) {
+	All(t, func(ctx *testcontext.Context, t *testing.T, db *metabase.DB) {
+		obj := randObjectStream()
+
+		t.Run("StreamID missing", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts:     metabase.ListStreamPositions{},
+				ErrClass: &metabase.ErrInvalidRequest,
+				ErrText:  "StreamID missing",
+			}.Check(ctx, t, db)
+
+			Verify{}.Check(ctx, t, db)
+		})
+
+		t.Run("Invalid limit", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    -1,
+				},
+				ErrClass: &metabase.ErrInvalidRequest,
+				ErrText:  "Invalid limit: -1",
+			}.Check(ctx, t, db)
+
+			Verify{}.Check(ctx, t, db)
+		})
+
+		t.Run("List no segments", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    1,
+				},
+				Result: metabase.ListStreamPositionsResult{},
+			}.Check(ctx, t, db)
+
+			Verify{}.Check(ctx, t, db)
+		})
+
+		t.Run("List segments", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			expectedObject := createObject(ctx, t, db, obj, 10)
+
+			expectedSegment := metabase.Segment{
+				StreamID: obj.StreamID,
+				Position: metabase.SegmentPosition{
+					Index: 0,
+				},
+				RootPieceID:       storj.PieceID{1},
+				EncryptedKey:      []byte{3},
+				EncryptedKeyNonce: []byte{4},
+				EncryptedSize:     1024,
+				PlainSize:         512,
+				Pieces:            metabase.Pieces{{Number: 0, StorageNode: storj.NodeID{2}}},
+				Redundancy:        defaultTestRedundancy,
+			}
+
+			expectedRawSegments := make([]metabase.RawSegment, 10)
+			expectedSegments := make([]metabase.SegmentPositionInfo, 10)
+			for i := range expectedSegments {
+				expectedSegment.Position.Index = uint32(i)
+				expectedRawSegments[i] = metabase.RawSegment(expectedSegment)
+				expectedSegments[i] = metabase.SegmentPositionInfo{
+					Position:  expectedSegment.Position,
+					PlainSize: expectedSegment.PlainSize,
+				}
+			}
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    10,
+				},
+				Result: metabase.ListStreamPositionsResult{
+					Segments: expectedSegments,
+				},
+			}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    1,
+				},
+				Result: metabase.ListStreamPositionsResult{
+					Segments: expectedSegments[:1],
+					More:     true,
+				},
+			}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    2,
+					Cursor: metabase.SegmentPosition{
+						Index: 1,
+					},
+				},
+				Result: metabase.ListStreamPositionsResult{
+					Segments: expectedSegments[2:4],
+					More:     true,
+				},
+			}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    2,
+					Cursor: metabase.SegmentPosition{
+						Index: 10,
+					},
+				},
+				Result: metabase.ListStreamPositionsResult{
+					More: false,
+				},
+			}.Check(ctx, t, db)
+
+			ListStreamPositions{
+				Opts: metabase.ListStreamPositions{
+					StreamID: obj.StreamID,
+					Limit:    2,
+					Cursor: metabase.SegmentPosition{
+						Part:  1,
+						Index: 10,
+					},
+				},
+				Result: metabase.ListStreamPositionsResult{
+					More: false,
+				},
+			}.Check(ctx, t, db)
+
+			Verify{
+				Objects: []metabase.RawObject{
+					metabase.RawObject(expectedObject),
+				},
+				Segments: expectedRawSegments,
+			}.Check(ctx, t, db)
+		})
+
+		t.Run("List segments from unordered parts", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			var testCases = []struct {
+				segments []metabase.SegmentPosition
+			}{
+				{[]metabase.SegmentPosition{
+					{Part: 3, Index: 0},
+					{Part: 0, Index: 0},
+					{Part: 1, Index: 0},
+					{Part: 2, Index: 0},
+				}},
+				{[]metabase.SegmentPosition{
+					{Part: 3, Index: 0},
+					{Part: 2, Index: 0},
+					{Part: 1, Index: 0},
+					{Part: 0, Index: 0},
+				}},
+				{[]metabase.SegmentPosition{
+					{Part: 0, Index: 0},
+					{Part: 2, Index: 0},
+					{Part: 1, Index: 0},
+					{Part: 3, Index: 0},
+				}},
+			}
+
+			expectedSegment := metabase.Segment{
+				StreamID:          obj.StreamID,
+				RootPieceID:       storj.PieceID{1},
+				EncryptedKey:      []byte{3},
+				EncryptedKeyNonce: []byte{4},
+				EncryptedSize:     1024,
+				PlainSize:         512,
+				Pieces:            metabase.Pieces{{Number: 0, StorageNode: storj.NodeID{2}}},
+				Redundancy:        defaultTestRedundancy,
+			}
+
+			for _, tc := range testCases {
+				obj := randObjectStream()
+
+				BeginObjectExactVersion{
+					Opts: metabase.BeginObjectExactVersion{
+						ObjectStream: obj,
+						Encryption:   defaultTestEncryption,
+					},
+					Version: obj.Version,
+				}.Check(ctx, t, db)
+
+				for i, segmentPosition := range tc.segments {
+					BeginSegment{
+						Opts: metabase.BeginSegment{
+							ObjectStream: obj,
+							Position:     segmentPosition,
+							RootPieceID:  storj.PieceID{byte(i + 1)},
+							Pieces: []metabase.Piece{{
+								Number:      1,
+								StorageNode: testrand.NodeID(),
+							}},
+						},
+					}.Check(ctx, t, db)
+
+					CommitSegment{
+						Opts: metabase.CommitSegment{
+							ObjectStream: obj,
+							Position:     segmentPosition,
+							RootPieceID:  storj.PieceID{1},
+							Pieces:       metabase.Pieces{{Number: 0, StorageNode: storj.NodeID{2}}},
+
+							EncryptedKey:      []byte{3},
+							EncryptedKeyNonce: []byte{4},
+
+							EncryptedSize: 1024,
+							PlainSize:     512,
+							PlainOffset:   0,
+							Redundancy:    defaultTestRedundancy,
+						},
+					}.Check(ctx, t, db)
+				}
+
+				CommitObject{
+					Opts: metabase.CommitObject{
+						ObjectStream: obj,
+					},
+				}.Check(ctx, t, db)
+
+				expectedSegments := make([]metabase.SegmentPositionInfo, 4)
+				for i := range expectedSegments {
+					pos := expectedSegment.Position
+					pos.Part = uint32(i)
+					expectedSegments[i] = metabase.SegmentPositionInfo{
+						Position:  pos,
+						PlainSize: expectedSegment.PlainSize,
+					}
+				}
+
+				ListStreamPositions{
+					Opts: metabase.ListStreamPositions{
+						StreamID: obj.StreamID,
+						Limit:    0,
+					},
+					Result: metabase.ListStreamPositionsResult{
+						Segments: expectedSegments,
+					},
+				}.Check(ctx, t, db)
+			}
+		})
+	})
+}

satellite/metainfo/metabase/test_test.go

@@ -274,6 +274,21 @@ func (step ListSegments) Check(ctx *testcontext.Context, t testing.TB, db *metab
 	require.Zero(t, diff)
 }
 
+type ListStreamPositions struct {
+	Opts     metabase.ListStreamPositions
+	Result   metabase.ListStreamPositionsResult
+	ErrClass *errs.Class
+	ErrText  string
+}
+
+func (step ListStreamPositions) Check(ctx *testcontext.Context, t testing.TB, db *metabase.DB) {
+	result, err := db.ListStreamPositions(ctx, step.Opts)
+	checkError(t, err, step.ErrClass, step.ErrText)
+
+	diff := cmp.Diff(step.Result, result, cmpopts.EquateApproxTime(5*time.Second))
+	require.Zero(t, diff)
+}
+
 type IterateLoopStreams struct {
 	Opts     metabase.IterateLoopStreams
 	Result   map[uuid.UUID][]metabase.LoopSegmentEntry