From ae78470b839064734be67989c31f2c10d7a829bb Mon Sep 17 00:00:00 2001
From: Patrick O'Grady <prohb125@gmail.com>
Date: Tue, 27 Oct 2020 09:58:56 -0700
Subject: [PATCH] First pass at tracking balance

---
 indexer/balance_storage_handler.go | 31 ++++++++++++++++++++
 indexer/balance_storage_helper.go  | 47 ++++++++++++++++++++++++++++++
 indexer/indexer.go                 |  9 +++++-
 3 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 indexer/balance_storage_handler.go
 create mode 100644 indexer/balance_storage_helper.go

diff --git a/indexer/balance_storage_handler.go b/indexer/balance_storage_handler.go
new file mode 100644
index 0000000..952753e
--- /dev/null
+++ b/indexer/balance_storage_handler.go
@@ -0,0 +1,31 @@
+package indexer
+
+import (
+	"context"
+
+	"github.com/coinbase/rosetta-sdk-go/parser"
+	"github.com/coinbase/rosetta-sdk-go/storage"
+	"github.com/coinbase/rosetta-sdk-go/types"
+)
+
+var _ storage.BalanceStorageHandler = (*BalanceStorageHandler)(nil)
+
+type BalanceStorageHandler struct{}
+
+// BlockAdded is called whenever a block is committed to BlockStorage.
+func (h *BalanceStorageHandler) BlockAdded(
+	ctx context.Context,
+	block *types.Block,
+	changes []*parser.BalanceChange,
+) error {
+	return nil
+}
+
+// BlockRemoved is called whenever a block is removed from BlockStorage.
+func (h *BalanceStorageHandler) BlockRemoved(
+	ctx context.Context,
+	block *types.Block,
+	changes []*parser.BalanceChange,
+) error {
+	return nil
+}
diff --git a/indexer/balance_storage_helper.go b/indexer/balance_storage_helper.go
new file mode 100644
index 0000000..fc1e0ca
--- /dev/null
+++ b/indexer/balance_storage_helper.go
@@ -0,0 +1,47 @@
+package indexer
+
+import (
+	"context"
+
+	"github.com/coinbase/rosetta-sdk-go/asserter"
+	"github.com/coinbase/rosetta-sdk-go/parser"
+	"github.com/coinbase/rosetta-sdk-go/storage"
+	"github.com/coinbase/rosetta-sdk-go/types"
+)
+
+var _ storage.BalanceStorageHelper = (*BalanceStorageHelper)(nil)
+
+type BalanceStorageHelper struct {
+	a *asserter.Asserter
+}
+
+// AccountBalance attempts to fetch the balance
+// for a missing account in storage.
+func (h *BalanceStorageHelper) AccountBalance(
+	ctx context.Context,
+	account *types.AccountIdentifier,
+	currency *types.Currency,
+	block *types.BlockIdentifier,
+) (*types.Amount, error) {
+	return &types.Amount{
+		Value:    "0",
+		Currency: currency,
+	}, nil
+}
+
+// Asserter returns a *asserter.Asserter.
+func (h *BalanceStorageHelper) Asserter() *asserter.Asserter {
+	return h.a
+}
+
+// BalanceExemptions returns a list of *types.BalanceExemption.
+func (h *BalanceStorageHelper) BalanceExemptions() []*types.BalanceExemption {
+	return []*types.BalanceExemption{}
+}
+
+// ExemptFunc returns a parser.ExemptOperation.
+func (h *BalanceStorageHelper) ExemptFunc() parser.ExemptOperation {
+	return func(op *types.Operation) bool {
+		return false
+	}
+}
diff --git a/indexer/indexer.go b/indexer/indexer.go
index 3f33757..3c7d04b 100644
--- a/indexer/indexer.go
+++ b/indexer/indexer.go
@@ -199,7 +199,14 @@ func Initialize(
 
 	coinStorage := storage.NewCoinStorage(localStore, i, asserter)
 	i.coinStorage = coinStorage
-	i.workers = []storage.BlockWorker{coinStorage}
+
+	balanceStorage := storage.NewBalanceStorage(localStore)
+	balanceStorage.Initialize(
+		&BalanceStorageHelper{asserter},
+		&BalanceStorageHandler{},
+	)
+
+	i.workers = []storage.BlockWorker{coinStorage, balanceStorage}
 
 	return i, nil
 }