storj/satellite/metabase/streamstat.go

66 lines
1.7 KiB
Go
Raw Normal View History

// Copyright (C) 2021 Storj Labs, Inc.
// See LICENSE for copying information.
package metabase
import (
"context"
"database/sql"
"errors"
"storj.io/common/storj"
"storj.io/common/uuid"
"storj.io/private/tagsql"
)
// GetStreamPieceCountByNodeID contains arguments for GetStreamPieceCountByNodeID.
type GetStreamPieceCountByNodeID struct {
StreamID uuid.UUID
}
// GetStreamPieceCountByNodeID returns piece count by node id.
func (db *DB) GetStreamPieceCountByNodeID(ctx context.Context, opts GetStreamPieceCountByNodeID) (result map[storj.NodeID]int64, err error) {
defer mon.Task()(&ctx)(&err)
if opts.StreamID.IsZero() {
return nil, ErrInvalidRequest.New("StreamID missing")
}
countByAlias := map[NodeAlias]int64{}
result = map[storj.NodeID]int64{}
err = withRows(db.db.QueryContext(ctx, `
SELECT remote_alias_pieces
FROM segments
WHERE stream_id IN (SELECT $1 UNION SELECT ancestor_stream_id FROM segment_copies WHERE stream_id = $1) AND remote_alias_pieces IS NOT null
`, opts.StreamID))(func(rows tagsql.Rows) error {
for rows.Next() {
var aliasPieces AliasPieces
err = rows.Scan(&aliasPieces)
if err != nil {
return Error.New("failed to scan segments: %w", err)
}
for i := range aliasPieces {
countByAlias[aliasPieces[i].Alias]++
}
}
return nil
})
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return result, nil
}
return result, Error.New("unable to fetch object segments: %w", err)
}
for alias, count := range countByAlias {
nodeID, err := db.aliasCache.Nodes(ctx, []NodeAlias{alias})
if err != nil {
return nil, Error.New("unable to convert aliases to pieces: %w", err)
}
result[nodeID[0]] = count
}
return result, nil
}