a2e20c93ae
Initially we duplicated the code to avoid large scale changes to the packages. Now we are past metainfo refactor we can remove the duplication. Change-Id: I9d0b2756cc6e2a2f4d576afa408a15273a7e1cef
66 lines
1.6 KiB
Go
66 lines
1.6 KiB
Go
// 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.Query(ctx, `
|
|
SELECT remote_alias_pieces
|
|
FROM segments
|
|
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
|
|
}
|