describe topic partition stats

Author: qrort

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `topicoptions.IncludePartitionStats()` for `Topic().Describe()` in order to get partition stats from server
+
 ## v3.106.1
 * Dropped `internal/allocator` package and all usages of it for further switch (test) protobuf opaque API
 

internal/grpcwrapper/rawtopic/describe_topic.go

@@ -17,12 +17,14 @@ import (
 type DescribeTopicRequest struct {
 	OperationParams rawydb.OperationParams
 	Path            string
+	IncludeStats    bool
 }
 
 func (req *DescribeTopicRequest) ToProto() *Ydb_Topic.DescribeTopicRequest {
 	return &Ydb_Topic.DescribeTopicRequest{
 		OperationParams: req.OperationParams.ToProto(),
 		Path:            req.Path,
+		IncludeStats:    req.IncludeStats,
 	}
 }
 
@@ -63,7 +65,10 @@ func (res *DescribeTopicResult) FromProto(response operation.Response) error {
 	protoPartitions := protoResult.GetPartitions()
 	res.Partitions = make([]PartitionInfo, len(protoPartitions))
 	for i, protoPartition := range protoPartitions {
-		res.Partitions[i].mustFromProto(protoPartition)
+		err := res.Partitions[i].FromProto(protoPartition)
+		if err != nil {
+			return err
+		}
 	}
 
 	res.RetentionPeriod = protoResult.GetRetentionPeriod().AsDuration()
@@ -93,12 +98,15 @@ type PartitionInfo struct {
 	Active             bool
 	ChildPartitionIDs  []int64
 	ParentPartitionIDs []int64
+	PartitionStats     PartitionStats
 }
 
-func (pi *PartitionInfo) mustFromProto(proto *Ydb_Topic.DescribeTopicResult_PartitionInfo) {
+func (pi *PartitionInfo) FromProto(proto *Ydb_Topic.DescribeTopicResult_PartitionInfo) error {
 	pi.PartitionID = proto.GetPartitionId()
 	pi.Active = proto.GetActive()
 
 	pi.ChildPartitionIDs = clone.Int64Slice(proto.GetChildPartitionIds())
 	pi.ParentPartitionIDs = clone.Int64Slice(proto.GetParentPartitionIds())
+
+	return pi.PartitionStats.FromProto(proto.GetPartitionStats())
 }

tests/integration/topic_client_test.go

@@ -4,6 +4,7 @@
 package integration
 
 import (
+	"bytes"
 	"context"
 	"io"
 	"os"
@@ -139,6 +140,95 @@ func TestTopicDescribe(t *testing.T) {
 	require.Equal(t, expected, topicDesc)
 }
 
+func TestTopicDescribePartitionStats(t *testing.T) {
+	ctx := xtest.Context(t)
+	db := connect(t)
+	topicName := "test-topic-" + t.Name()
+
+	var (
+		supportedCodecs     = []topictypes.Codec{topictypes.CodecRaw}
+		minActivePartitions = int64(1)
+	)
+
+	_ = db.Topic().Drop(ctx, topicName)
+	err := db.Topic().Create(ctx, topicName,
+		topicoptions.CreateWithSupportedCodecs(supportedCodecs...),
+		topicoptions.CreateWithMinActivePartitions(minActivePartitions),
+	)
+	require.NoError(t, err)
+
+	writer, err := db.Topic().StartWriter(topicName, topicoptions.WithWriterWaitServerAck(true))
+	require.NoError(t, err)
+	topicMessage := topicwriter.Message{
+		Data: bytes.NewReader([]byte{128}),
+	}
+	err = writer.Write(ctx, topicMessage)
+	require.NoError(t, err)
+
+	topicDesc, err := db.Topic().Describe(ctx, topicName, topicoptions.IncludePartitionStats())
+	require.NoError(t, err)
+
+	expected := topictypes.TopicDescription{
+		Path: topicName,
+		PartitionSettings: topictypes.PartitionSettings{
+			MinActivePartitions: minActivePartitions,
+		},
+		Partitions: []topictypes.PartitionInfo{
+			{
+				PartitionID: 0,
+				Active:      true,
+				PartitionStats: topictypes.PartitionStats{
+					PartitionsOffset: topictypes.OffsetRange{
+						Start: 0,
+						End:   1,
+					},
+					BytesWritten: topictypes.MultipleWindowsStat{
+						PerMinute: 1,
+						PerHour:   1,
+						PerDay:    1,
+					},
+				},
+			},
+		},
+		RetentionPeriod:                   86400000000000,
+		RetentionStorageMB:                0,
+		SupportedCodecs:                   supportedCodecs,
+		PartitionWriteBurstBytes:          1048576,
+		PartitionWriteSpeedBytesPerSecond: 1048576,
+		Attributes:                        nil,
+		Consumers:                         []topictypes.Consumer{},
+		MeteringMode:                      topictypes.MeteringModeUnspecified,
+	}
+
+	requireAndCleanSubset := func(checked *map[string]string, subset *map[string]string) {
+		t.Helper()
+		for k, subValue := range *subset {
+			checkedValue, ok := (*checked)[k]
+			require.True(t, ok, k)
+			require.Equal(t, subValue, checkedValue)
+		}
+		*checked = nil
+		*subset = nil
+	}
+
+	topicDesc.Partitions[0].PartitionStats.LastWriteTime = nil
+	bw := &topicDesc.Partitions[0].PartitionStats.BytesWritten
+	sign := func(x int64) int64 {
+		if x > 0 {
+			return 1
+		}
+		return 0
+	}
+	topicDesc.Partitions[0].PartitionStats.BytesWritten = topictypes.MultipleWindowsStat{
+		PerMinute: sign(bw.PerMinute),
+		PerHour:   sign(bw.PerHour),
+		PerDay:    sign(bw.PerDay),
+	}
+
+	requireAndCleanSubset(&topicDesc.Attributes, &expected.Attributes)
+	require.Equal(t, expected, topicDesc)
+}
+
 func TestDescribeTopicConsumer(t *testing.T) {
 	ctx := xtest.Context(t)
 	db := connect(t)

topic/topicoptions/topicoptions_describe.go

@@ -2,12 +2,21 @@ package topicoptions
 
 import "github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic"
 
-// DescribeOption type for options of describe method. Not used now.
+// DescribeOption type for options of describe method.
 type DescribeOption func(req *rawtopic.DescribeTopicRequest)
 
+// IncludePartitionStats additionally fetches DescribeTopicResult.PartitionInfo.PartitionSettings from server
+func IncludePartitionStats() DescribeOption {
+	return func(req *rawtopic.DescribeTopicRequest) {
+		req.IncludeStats = true
+	}
+}
+
 // DescribeConsumerOption type for options of describe consumer method.
 type DescribeConsumerOption func(req *rawtopic.DescribeConsumerRequest)
 
+// IncludeConsumerStats additionally fetches
+// DescribeConsumerResult.DescribeConsumerResultPartitionInfo.PartitionConsumerStats from server
 func IncludeConsumerStats() DescribeConsumerOption {
 	return func(req *rawtopic.DescribeConsumerRequest) {
 		req.IncludeStats = true

topic/topictypes/topictypes.go

@@ -169,6 +169,7 @@ type PartitionInfo struct {
 	Active             bool
 	ChildPartitionIDs  []int64
 	ParentPartitionIDs []int64
+	PartitionStats     PartitionStats
 }
 
 // FromRaw convert from internal format to public. Used internally only.
@@ -178,6 +179,7 @@ func (p *PartitionInfo) FromRaw(raw *rawtopic.PartitionInfo) {
 
 	p.ChildPartitionIDs = clone.Int64Slice(raw.ChildPartitionIDs)
 	p.ParentPartitionIDs = clone.Int64Slice(raw.ParentPartitionIDs)
+	p.PartitionStats.FromRaw(&raw.PartitionStats)
 }
 
 type MultipleWindowsStat struct {

topic/topictypes/topictypes_test.go

@@ -11,6 +11,11 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
 )
 
+var (
+	fourPM = time.Date(2024, 0o1, 0o1, 16, 0, 0, 0, time.UTC)
+	hour   = time.Hour
+)
+
 func TestTopicDescriptionFromRaw(t *testing.T) {
 	testData := []struct {
 		testName            string
@@ -64,6 +69,20 @@ func TestTopicDescriptionFromRaw(t *testing.T) {
 						ParentPartitionIDs: []int64{
 							1, 2, 3,
 						},
+						PartitionStats: PartitionStats{
+							PartitionsOffset: OffsetRange{
+								Start: 10,
+								End:   20,
+							},
+							StoreSizeBytes:  1024,
+							LastWriteTime:   &fourPM,
+							MaxWriteTimeLag: &hour,
+							BytesWritten: MultipleWindowsStat{
+								PerMinute: 1,
+								PerHour:   60,
+								PerDay:    1440,
+							},
+						},
 					},
 					{
 						PartitionID: 43,
@@ -104,6 +123,26 @@ func TestTopicDescriptionFromRaw(t *testing.T) {
 						ParentPartitionIDs: []int64{
 							1, 2, 3,
 						},
+						PartitionStats: rawtopic.PartitionStats{
+							PartitionsOffset: rawtopiccommon.OffsetRange{
+								Start: 10,
+								End:   20,
+							},
+							StoreSizeBytes: 1024,
+							LastWriteTime: rawoptional.Time{
+								Value:    fourPM,
+								HasValue: true,
+							},
+							MaxWriteTimeLag: rawoptional.Duration{
+								Value:    hour,
+								HasValue: true,
+							},
+							BytesWritten: rawtopic.MultipleWindowsStat{
+								PerMinute: 1,
+								PerHour:   60,
+								PerDay:    1440,
+							},
+						},
 					},
 					{
 						PartitionID: 43,
@@ -165,8 +204,6 @@ func TestTopicDescriptionFromRaw(t *testing.T) {
 }
 
 func TestTopicConsumerDescriptionFromRaw(t *testing.T) {
-	fourPM := time.Date(2024, 0o1, 0o1, 16, 0, 0, 0, time.UTC)
-	hour := time.Hour
 	testData := []struct {
 		testName               string
 		expectedDescription    TopicConsumerDescription