Add ledger client dataprovider.

Makefile

@@ -19,6 +19,7 @@ test-cover:
 gen:
 	cd candid && go generate
 	cd pocketic && go generate
+	cd clients/ledger && go generate
 	cd clients/registry && go generate
 
 gen-ic:

clients/ledger/client_test.go

@@ -0,0 +1,13 @@
+package ledger_test
+
+import (
+	"os"
+	"testing"
+)
+
+func checkEnabled(t *testing.T) {
+	// The reason for this is that the tests are very slow.
+	if os.Getenv("LEDGER_TEST_ENABLE") != "true" {
+		t.Skip("Skipping registry tests. Set LEDGER_TEST_ENABLE=true to enable.")
+	}
+}

clients/ledger/dataprovider.go

@@ -0,0 +1,139 @@
+package ledger
+
+import (
+	"fmt"
+	"github.com/aviate-labs/agent-go"
+	v1 "github.com/aviate-labs/agent-go/clients/ledger/proto/v1"
+	"github.com/aviate-labs/agent-go/ic"
+	"github.com/aviate-labs/agent-go/principal"
+)
+
+const MaxBlocksPerRequest = 2000
+
+type BlockIndex uint64
+
+type DataProvider struct {
+	a *agent.Agent
+}
+
+func NewDataProvider() (*DataProvider, error) {
+	a, err := agent.New(agent.DefaultConfig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create agent: %w", err)
+	}
+	return &DataProvider{a: a}, nil
+}
+
+func (d DataProvider) GetArchiveIndex() ([]*v1.ArchiveIndexEntry, error) {
+	var resp v1.ArchiveIndexResponse
+	if err := d.a.QueryProto(
+		ic.LEDGER_PRINCIPAL,
+		"get_archive_index_pb",
+		nil,
+		&resp,
+	); err != nil {
+		return nil, fmt.Errorf("failed to get archive index: %w", err)
+	}
+	return resp.Entries, nil
+}
+
+func (d DataProvider) GetRawBlock(height BlockIndex) (*v1.EncodedBlock, error) {
+	var resp v1.BlockResponse
+	if err := d.a.QueryProto(
+		ic.LEDGER_PRINCIPAL,
+		"block_pb",
+		&v1.BlockRequest{
+			BlockHeight: uint64(height),
+		},
+		&resp,
+	); err != nil {
+		return nil, fmt.Errorf("failed to get block: %w", err)
+	}
+	switch blockResponse := resp.BlockContent.(type) {
+	case *v1.BlockResponse_Block:
+		return blockResponse.Block, nil
+	case *v1.BlockResponse_CanisterId:
+		archiveCanisterID := principal.Principal{Raw: blockResponse.CanisterId.SerializedId}
+		var archiveResp v1.BlockResponse
+		if err := d.a.QueryProto(
+			archiveCanisterID,
+			"get_block_pb",
+			&v1.BlockRequest{
+				BlockHeight: uint64(height),
+			},
+			&archiveResp,
+		); err != nil {
+			return nil, fmt.Errorf("failed to get blocks: %w", err)
+		}
+		// Will never return a CanisterId block.
+		return archiveResp.GetBlock(), nil
+	default:
+		return nil, fmt.Errorf("unexpected block content type: %T", blockResponse)
+	}
+}
+
+func (d DataProvider) GetRawBlocks(start, end BlockIndex) ([]*v1.EncodedBlock, error) {
+	if end-start < 2000 {
+		blocks, err := d.GetRawBlocksRange(ic.LEDGER_PRINCIPAL, start, end)
+		if err == nil {
+			return blocks, nil
+		}
+	}
+	archives, err := d.GetArchiveIndex()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get archive index: %w", err)
+	}
+	var blocks []*v1.EncodedBlock
+	for _, archive := range archives {
+		if archive.HeightTo < uint64(start) || uint64(end) < archive.HeightFrom {
+			continue
+		}
+		for start < min(BlockIndex(archive.HeightTo), end) {
+			archiveEnd := min(end, BlockIndex(archive.HeightTo), start+MaxBlocksPerRequest)
+			archiveBlocks, err := d.GetRawBlocksRange(principal.Principal{Raw: archive.CanisterId.SerializedId}, start, archiveEnd)
+			if err != nil {
+				return nil, fmt.Errorf("failed to get archive blocks: %w", err)
+			}
+			blocks = append(blocks, archiveBlocks...)
+			start += BlockIndex(len(archiveBlocks))
+		}
+	}
+	return blocks, nil
+}
+
+func (d DataProvider) GetRawBlocksRange(canisterID principal.Principal, start, end BlockIndex) ([]*v1.EncodedBlock, error) {
+	var resp v1.GetBlocksResponse
+	if err := d.a.QueryProto(
+		canisterID,
+		"get_blocks_pb",
+		&v1.GetBlocksRequest{
+			Start:  uint64(start),
+			Length: uint64(end - start),
+		},
+		&resp,
+	); err != nil {
+		return nil, fmt.Errorf("failed to get blocks: %w", err)
+	}
+	switch blocksResponse := resp.GetBlocksContent.(type) {
+	case *v1.GetBlocksResponse_Blocks:
+		return blocksResponse.Blocks.Blocks, nil
+	case *v1.GetBlocksResponse_Error:
+		return nil, fmt.Errorf("failed to get blocks: %s", blocksResponse.Error)
+	default:
+		return nil, fmt.Errorf("unexpected get block content type: %T", blocksResponse)
+	}
+}
+
+func (d DataProvider) GetTipOfChain() (*BlockIndex, error) {
+	var resp v1.TipOfChainResponse
+	if err := d.a.QueryProto(
+		ic.LEDGER_PRINCIPAL,
+		"tip_of_chain_pb",
+		&v1.TipOfChainRequest{},
+		&resp,
+	); err != nil {
+		return nil, fmt.Errorf("failed to get tip of chain: %w", err)
+	}
+	height := BlockIndex(resp.ChainLength.Height)
+	return &height, nil
+}

clients/ledger/dataprovider_test.go

@@ -0,0 +1,47 @@
+package ledger_test
+
+import (
+	"github.com/aviate-labs/agent-go/clients/ledger"
+	"testing"
+)
+
+func TestDataProvider_GetRawBlock(t *testing.T) {
+	checkEnabled(t)
+
+	dp, err := ledger.NewDataProvider()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := dp.GetRawBlock(0); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestDataProvider_GetRawBlocks(t *testing.T) {
+	checkEnabled(t)
+
+	dp, err := ledger.NewDataProvider()
+	if err != nil {
+		t.Fatal(err)
+	}
+	n := 3 * ledger.MaxBlocksPerRequest
+	blocks, err := dp.GetRawBlocks(0, ledger.BlockIndex(n))
+	if err != nil {
+		t.Error(err)
+	}
+	if len(blocks) != n {
+		t.Errorf("expected %d blocks, got %d", n, len(blocks))
+	}
+}
+
+func TestDataProvider_GetTipOfChain(t *testing.T) {
+	checkEnabled(t)
+
+	dp, err := ledger.NewDataProvider()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := dp.GetTipOfChain(); err != nil {
+		t.Error(err)
+	}
+}

clients/ledger/proto.go

@@ -0,0 +1,4 @@
+package ledger
+
+//go:generate go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
+//go:generate protoc -I=testdata --go_out=. testdata/ledger.proto