Use page token when list cluster metadata from DB (#2331)

temporalio · Dec 29, 2021 · 1e02664 · 1e02664
1 parent e9a1cc9
commit 1e02664
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 3 deletions.
diff --git a/common/cluster/metadata.go b/common/cluster/metadata.go
@@ -477,7 +477,10 @@ func (m *metadataImpl) listAllClusterMetadataFromDB() (map[string]*ClusterInform
 
 	paginationFn := func(paginationToken []byte) ([]interface{}, []byte, error) {
 		resp, err := m.clusterMetadataStore.ListClusterMetadata(
-			&persistence.ListClusterMetadataRequest{PageSize: defaultClusterMetadataPageSize},
+			&persistence.ListClusterMetadataRequest{
+				PageSize:      defaultClusterMetadataPageSize,
+				NextPageToken: paginationToken,
+			},
 		)
 		if err != nil {
 			return nil, nil, err

diff --git a/common/cluster/metadata_test.go b/common/cluster/metadata_test.go
@@ -138,7 +138,7 @@ func (s *metadataSuite) Test_RegisterMetadataChangeCallback() {
 	s.Equal(0, len(s.metadata.clusterChangeCallback))
 }
 
-func (s *metadataSuite) Test_RefreshClusterMetadata() {
+func (s *metadataSuite) Test_RefreshClusterMetadata_Success() {
 	id := uuid.New()
 	s.metadata.clusterChangeCallback[id] = func(oldClusterMetadata map[string]*ClusterInformation, newClusterMetadata map[string]*ClusterInformation) {
 		oldMetadata, ok := oldClusterMetadata[id]
@@ -179,5 +179,50 @@ func (s *metadataSuite) Test_RefreshClusterMetadata() {
 				},
 			},
 		}, nil)
-	s.metadata.refreshClusterMetadata(context.Background())
+	err := s.metadata.refreshClusterMetadata(context.Background())
+	s.NoError(err)
+}
+
+func (s *metadataSuite) Test_ListAllClusterMetadataFromDB_Success() {
+	nextPageSizeToken := []byte{1}
+	newClusterName := uuid.New()
+	s.mockClusterMetadataStore.EXPECT().ListClusterMetadata(&persistence.ListClusterMetadataRequest{
+		PageSize:      defaultClusterMetadataPageSize,
+		NextPageToken: nil,
+	}).Return(
+		&persistence.ListClusterMetadataResponse{
+			ClusterMetadata: []*persistence.GetClusterMetadataResponse{
+				{
+					ClusterMetadata: persistencespb.ClusterMetadata{
+						ClusterName:            s.clusterName,
+						IsConnectionEnabled:    true,
+						InitialFailoverVersion: 1,
+						ClusterAddress:         uuid.New(),
+					},
+					Version: 1,
+				},
+			},
+			NextPageToken: nextPageSizeToken,
+		}, nil).Times(1)
+	s.mockClusterMetadataStore.EXPECT().ListClusterMetadata(&persistence.ListClusterMetadataRequest{
+		PageSize:      defaultClusterMetadataPageSize,
+		NextPageToken: nextPageSizeToken,
+	}).Return(
+		&persistence.ListClusterMetadataResponse{
+			ClusterMetadata: []*persistence.GetClusterMetadataResponse{
+				{
+					ClusterMetadata: persistencespb.ClusterMetadata{
+						ClusterName:            newClusterName,
+						IsConnectionEnabled:    true,
+						InitialFailoverVersion: 2,
+						ClusterAddress:         uuid.New(),
+					},
+					Version: 2,
+				},
+			},
+		}, nil).Times(1)
+
+	resp, err := s.metadata.listAllClusterMetadataFromDB()
+	s.NoError(err)
+	s.Equal(2, len(resp))
 }