temporalio · MichaelSnowden · Sep 30, 2022 · Sep 27, 2022
diff --git a/common/persistence/persistence-tests/historyV2PersistenceTest.go b/common/persistence/persistence-tests/historyV2PersistenceTest.go
@@ -122,11 +122,6 @@ func (s *HistoryV2PersistenceSuite) TestGenUUIDs() {
 
 // TestScanAllTrees test
 func (s *HistoryV2PersistenceSuite) TestScanAllTrees() {
-	// TODO https://go.temporal.io/server/issues/2458
-	if s.ExecutionManager.GetName() != "cassandra" {
-		return
-	}
-
 	resp, err := s.ExecutionManager.GetAllHistoryTreeBranches(s.ctx, &p.GetAllHistoryTreeBranchesRequest{
 		PageSize: 1,
 	})

diff --git a/common/persistence/persistenceInterface.go b/common/persistence/persistenceInterface.go
@@ -155,7 +155,9 @@ type (
 		DeleteHistoryBranch(ctx context.Context, request *InternalDeleteHistoryBranchRequest) error
 		// GetHistoryTree returns all branch information of a tree
 		GetHistoryTree(ctx context.Context, request *GetHistoryTreeRequest) (*InternalGetHistoryTreeResponse, error)
-		// GetAllHistoryTreeBranches returns all branches of all trees
+		// GetAllHistoryTreeBranches returns all branches of all trees.
+		// Note that branches may be skipped or duplicated across pages if there are branches created or deleted while
+		// paginating through results.
 		GetAllHistoryTreeBranches(ctx context.Context, request *GetAllHistoryTreeBranchesRequest) (*InternalGetAllHistoryTreeBranchesResponse, error)
 	}
 

diff --git a/common/persistence/sql/history_store.go b/common/persistence/sql/history_store.go
@@ -27,6 +27,7 @@ package sql
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"fmt"
 	"math"
 
@@ -425,14 +426,72 @@ func (m *sqlExecutionStore) DeleteHistoryBranch(
 	})
 }
 
+// getAllHistoryTreeBranchesPaginationToken represents the primary key of the latest row in the history_tree table that
+// we returned.
+type getAllHistoryTreeBranchesPaginationToken struct {
+	ShardID  int32
+	TreeID   primitives.UUID
+	BranchID primitives.UUID
+}
+
 func (m *sqlExecutionStore) GetAllHistoryTreeBranches(
-    	_ context.Context,
+	ctx context.Context,
 	request *p.GetAllHistoryTreeBranchesRequest,
 ) (*p.InternalGetAllHistoryTreeBranchesResponse, error) {
+	pageSize := request.PageSize
+	if pageSize <= 0 {
+		return nil, fmt.Errorf("PageSize must be greater than 0, but was %d", pageSize)
+	}
+
+	page := sqlplugin.HistoryTreeBranchPage{
+		Limit: pageSize,
+	}
+	if len(request.NextPageToken) != 0 {
+		var token getAllHistoryTreeBranchesPaginationToken
+		if err := json.Unmarshal(request.NextPageToken, &token); err != nil {
+			return nil, err
+		}
+		page.ShardID = token.ShardID
+		page.TreeID = token.TreeID
+		page.BranchID = token.BranchID
+	}
 
-	// TODO https://github.com/uber/cadence/issues/2458
-	// Implement it when we need
-	panic("not implemented yet")
+	rows, err := m.Db.PaginateBranchesFromHistoryTree(ctx, page)
+	if err != nil {
+		return nil, err
+	}
+	branches := make([]p.InternalHistoryBranchDetail, 0, pageSize)
+	for _, row := range rows {
+		branch := p.InternalHistoryBranchDetail{
+			TreeID:   row.TreeID.String(),
+			BranchID: row.BranchID.String(),
+			Data:     row.Data,
+			Encoding: row.DataEncoding,
+		}
+		branches = append(branches, branch)
+	}
+
+	response := &p.InternalGetAllHistoryTreeBranchesResponse{
+		Branches: branches,
+	}
+	if len(branches) < pageSize {
+		// no next page token because there are no more results
+		return response, nil
+	}
+
+	// if we filled the page with rows, then set the next page token
+	lastRow := rows[len(rows)-1]
+	token := getAllHistoryTreeBranchesPaginationToken{
+		ShardID:  lastRow.ShardID,
+		TreeID:   lastRow.TreeID,
+		BranchID: lastRow.BranchID,
+	}
+	tokenBytes, err := json.Marshal(token)
+	if err != nil {
+		return nil, err
+	}
+	response.NextPageToken = tokenBytes
+	return response, nil
 }
 
 // GetHistoryTree returns all branch information of a tree

diff --git a/common/persistence/sql/sqlplugin/history_tree.go b/common/persistence/sql/sqlplugin/history_tree.go
@@ -56,10 +56,19 @@ type (
 		BranchID primitives.UUID
 	}
 
-	// HistoryNode is the SQL persistence interface for history trees
+	// HistoryTreeBranchPage is a struct which represents a page of history tree branches to query.
+	HistoryTreeBranchPage struct {
+		ShardID  int32
+		TreeID   primitives.UUID
+		BranchID primitives.UUID
+		Limit    int
+	}
+
+	// HistoryTree is the SQL persistence interface for history trees
 	HistoryTree interface {
 		InsertIntoHistoryTree(ctx context.Context, row *HistoryTreeRow) (sql.Result, error)
 		SelectFromHistoryTree(ctx context.Context, filter HistoryTreeSelectFilter) ([]HistoryTreeRow, error)
 		DeleteFromHistoryTree(ctx context.Context, filter HistoryTreeDeleteFilter) (sql.Result, error)
+		PaginateBranchesFromHistoryTree(ctx context.Context, filter HistoryTreeBranchPage) ([]HistoryTreeRow, error)
 	}
 )
diff --git a/common/persistence/sql/sqlplugin/mysql/events.go b/common/persistence/sql/sqlplugin/mysql/events.go
@@ -62,6 +62,12 @@ const (
 
 	getHistoryTreeQuery = `SELECT branch_id, data, data_encoding FROM history_tree WHERE shard_id = ? AND tree_id = ? `
 
+	paginateBranchesQuery = `SELECT shard_id, tree_id, branch_id, data, data_encoding
+        FROM history_tree
+        WHERE (shard_id, tree_id, branch_id) > (?, ?, ?)
+        ORDER BY shard_id, tree_id, branch_id
+        LIMIT ?`
+
 	deleteHistoryTreeQuery = `DELETE FROM history_tree WHERE shard_id = ? AND tree_id = ? AND branch_id = ? `
 )
 
@@ -190,6 +196,25 @@ func (mdb *db) SelectFromHistoryTree(
 	return rows, err
 }
 
+// PaginateBranchesFromHistoryTree reads up to page.Limit rows from the history_tree table sorted by their primary key,
+// while skipping the first page.Offset rows.
+func (mdb *db) PaginateBranchesFromHistoryTree(
+	ctx context.Context,
+	page sqlplugin.HistoryTreeBranchPage,
+) ([]sqlplugin.HistoryTreeRow, error) {
+	var rows []sqlplugin.HistoryTreeRow
+	err := mdb.conn.SelectContext(
+		ctx,
+		&rows,
+		paginateBranchesQuery,
+		page.ShardID,
+		page.TreeID,
+		page.BranchID,
+		page.Limit,
+	)
+	return rows, err
+}
+
 // DeleteFromHistoryTree deletes one or more rows from history_tree table
 func (mdb *db) DeleteFromHistoryTree(
 	ctx context.Context,

diff --git a/common/persistence/sql/sqlplugin/postgresql/events.go b/common/persistence/sql/sqlplugin/postgresql/events.go
@@ -64,6 +64,12 @@ const (
 
 	getHistoryTreeQuery = `SELECT branch_id, data, data_encoding FROM history_tree WHERE shard_id = $1 AND tree_id = $2 `
 
+	paginateBranchesQuery = `SELECT shard_id, tree_id, branch_id, data, data_encoding
+        FROM history_tree
+        WHERE (shard_id, tree_id, branch_id) > ($1, $2, $3)
+        ORDER BY shard_id, tree_id, branch_id
+        LIMIT $4`
+
 	deleteHistoryTreeQuery = `DELETE FROM history_tree WHERE shard_id = $1 AND tree_id = $2 AND branch_id = $3 `
 )
 
@@ -192,6 +198,24 @@ func (pdb *db) SelectFromHistoryTree(
 	return rows, err
 }
 
+// PaginateBranchesFromHistoryTree reads up to page.Limit rows from the history_tree table sorted by their primary key,
+// while skipping the first page.Offset rows.
+func (pdb *db) PaginateBranchesFromHistoryTree(
+	ctx context.Context,
+	page sqlplugin.HistoryTreeBranchPage,
+) ([]sqlplugin.HistoryTreeRow, error) {
+	var rows []sqlplugin.HistoryTreeRow
+	err := pdb.conn.SelectContext(ctx,
+		&rows,
+		paginateBranchesQuery,
+		page.ShardID,
+		page.TreeID,
+		page.BranchID,
+		page.Limit,
+	)
+	return rows, err
+}
+
 // DeleteFromHistoryTree deletes one or more rows from history_tree table
 func (pdb *db) DeleteFromHistoryTree(
 	ctx context.Context,

diff --git a/common/persistence/sql/sqlplugin/sqlite/events.go b/common/persistence/sql/sqlplugin/sqlite/events.go
@@ -62,6 +62,12 @@ const (
 
 	getHistoryTreeQuery = `SELECT branch_id, data, data_encoding FROM history_tree WHERE shard_id = ? AND tree_id = ? `
 
+	paginateBranchesQuery = `SELECT shard_id, tree_id, branch_id, data, data_encoding
+        FROM history_tree
+        WHERE (shard_id, tree_id, branch_id) > ($1, $2, $3)
+        ORDER BY shard_id, tree_id, branch_id
+        LIMIT $4`
+
 	deleteHistoryTreeQuery = `DELETE FROM history_tree WHERE shard_id = ? AND tree_id = ? AND branch_id = ? `
 )
 
@@ -190,6 +196,24 @@ func (mdb *db) SelectFromHistoryTree(
 	return rows, err
 }
 
+// PaginateBranchesFromHistoryTree reads up to page.Limit rows from the history_tree table sorted by their primary key,
+// while skipping the first page.Offset rows.
+func (mdb *db) PaginateBranchesFromHistoryTree(
+	ctx context.Context,
+	page sqlplugin.HistoryTreeBranchPage,
+) ([]sqlplugin.HistoryTreeRow, error) {
+	var rows []sqlplugin.HistoryTreeRow
+	err := mdb.conn.SelectContext(ctx,
+		&rows,
+		paginateBranchesQuery,
+		page.ShardID,
+		page.TreeID,
+		page.BranchID,
+		page.Limit,
+	)
+	return rows, err
+}
+
 // DeleteFromHistoryTree deletes one or more rows from history_tree table
 func (mdb *db) DeleteFromHistoryTree(
 	ctx context.Context,