From 1ad5ccc50f60add7cd63ef048ce0638aaeeaf56a Mon Sep 17 00:00:00 2001
From: chyezh <chyezh@outlook.com>
Date: Mon, 1 Apr 2024 10:21:21 +0800
Subject: [PATCH] enhance: add rg and db interface for segment and db/rg metric
 label (#31715)

issue: #30931

Signed-off-by: chyezh <chyezh@outlook.com>
---
 internal/querynodev2/segments/mock_segment.go | 82 +++++++++++++++++++
 internal/querynodev2/segments/segment.go      |  8 ++
 .../querynodev2/segments/segment_interface.go |  2 +
 pkg/metrics/metrics.go                        |  2 +
 4 files changed, 94 insertions(+)

diff --git a/internal/querynodev2/segments/mock_segment.go b/internal/querynodev2/segments/mock_segment.go
index 874c61dddb..a433973321 100644
--- a/internal/querynodev2/segments/mock_segment.go
+++ b/internal/querynodev2/segments/mock_segment.go
@@ -117,6 +117,47 @@ func (_c *MockSegment_Collection_Call) RunAndReturn(run func() int64) *MockSegme
 	return _c
 }
 
+// DatabaseName provides a mock function with given fields:
+func (_m *MockSegment) DatabaseName() string {
+	ret := _m.Called()
+
+	var r0 string
+	if rf, ok := ret.Get(0).(func() string); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(string)
+	}
+
+	return r0
+}
+
+// MockSegment_DatabaseName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DatabaseName'
+type MockSegment_DatabaseName_Call struct {
+	*mock.Call
+}
+
+// DatabaseName is a helper method to define mock.On call
+func (_e *MockSegment_Expecter) DatabaseName() *MockSegment_DatabaseName_Call {
+	return &MockSegment_DatabaseName_Call{Call: _e.mock.On("DatabaseName")}
+}
+
+func (_c *MockSegment_DatabaseName_Call) Run(run func()) *MockSegment_DatabaseName_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockSegment_DatabaseName_Call) Return(_a0 string) *MockSegment_DatabaseName_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockSegment_DatabaseName_Call) RunAndReturn(run func() string) *MockSegment_DatabaseName_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Delete provides a mock function with given fields: ctx, primaryKeys, timestamps
 func (_m *MockSegment) Delete(ctx context.Context, primaryKeys []storage.PrimaryKey, timestamps []uint64) error {
 	ret := _m.Called(ctx, primaryKeys, timestamps)
@@ -952,6 +993,47 @@ func (_c *MockSegment_Release_Call) RunAndReturn(run func(...releaseOption)) *Mo
 	return _c
 }
 
+// ResourceGroup provides a mock function with given fields:
+func (_m *MockSegment) ResourceGroup() string {
+	ret := _m.Called()
+
+	var r0 string
+	if rf, ok := ret.Get(0).(func() string); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(string)
+	}
+
+	return r0
+}
+
+// MockSegment_ResourceGroup_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ResourceGroup'
+type MockSegment_ResourceGroup_Call struct {
+	*mock.Call
+}
+
+// ResourceGroup is a helper method to define mock.On call
+func (_e *MockSegment_Expecter) ResourceGroup() *MockSegment_ResourceGroup_Call {
+	return &MockSegment_ResourceGroup_Call{Call: _e.mock.On("ResourceGroup")}
+}
+
+func (_c *MockSegment_ResourceGroup_Call) Run(run func()) *MockSegment_ResourceGroup_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockSegment_ResourceGroup_Call) Return(_a0 string) *MockSegment_ResourceGroup_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockSegment_ResourceGroup_Call) RunAndReturn(run func() string) *MockSegment_ResourceGroup_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // ResourceUsageEstimate provides a mock function with given fields:
 func (_m *MockSegment) ResourceUsageEstimate() ResourceUsage {
 	ret := _m.Called()
diff --git a/internal/querynodev2/segments/segment.go b/internal/querynodev2/segments/segment.go
index 40775e3529..a602fa0d6e 100644
--- a/internal/querynodev2/segments/segment.go
+++ b/internal/querynodev2/segments/segment.go
@@ -120,6 +120,14 @@ func (s *baseSegment) Partition() int64 {
 	return s.loadInfo.GetPartitionID()
 }
 
+func (s *baseSegment) DatabaseName() string {
+	return s.collection.GetDBName()
+}
+
+func (s *baseSegment) ResourceGroup() string {
+	return s.collection.GetResourceGroup()
+}
+
 func (s *baseSegment) Shard() string {
 	return s.loadInfo.GetInsertChannel()
 }
diff --git a/internal/querynodev2/segments/segment_interface.go b/internal/querynodev2/segments/segment_interface.go
index 4e9b07a78f..0f9f1e37c8 100644
--- a/internal/querynodev2/segments/segment_interface.go
+++ b/internal/querynodev2/segments/segment_interface.go
@@ -50,6 +50,8 @@ type Segment interface {
 
 	// Properties
 	ID() int64
+	DatabaseName() string
+	ResourceGroup() string
 	Collection() int64
 	Partition() int64
 	Shard() string
diff --git a/pkg/metrics/metrics.go b/pkg/metrics/metrics.go
index edc2b9c475..4ba5adfb34 100644
--- a/pkg/metrics/metrics.go
+++ b/pkg/metrics/metrics.go
@@ -81,6 +81,8 @@ const (
 	functionLabelName        = "function_name"
 	queryTypeLabelName       = "query_type"
 	collectionName           = "collection_name"
+	databaseLabelName        = "db_name"
+	resourceGroupLabelName   = "rg"
 	indexName                = "index_name"
 	isVectorIndex            = "is_vector_index"
 	segmentStateLabelName    = "segment_state"