From e9ca581becbb6b6ffd6f7759763841cb36fd0297 Mon Sep 17 00:00:00 2001
From: Dave Collins <davec@conformal.com>
Date: Wed, 21 Oct 2020 13:13:03 -0500
Subject: [PATCH] blockchain: Create treasury buckets during upgrade.

This modifies the database upgrade code in blockchain to properly create
the buckets needed to house the decentralized treasury information.

It is done in a such a way that it will correct any upgrades that have
already taken place before this code is active.
---
 blockchain/chainio.go |  4 ++--
 blockchain/upgrade.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/blockchain/chainio.go b/blockchain/chainio.go
index d6d52281..fc690d33 100644
--- a/blockchain/chainio.go
+++ b/blockchain/chainio.go
@@ -1695,8 +1695,8 @@ func (b *BlockChain) createChainState() error {
 			return err
 		}
 
-		// Create the bucket that houses the treasury account
-		// information.
+		// Create the buckets that house the treasury account and spend
+		// transaction information.
 		_, err = meta.CreateBucket(treasuryBucketName)
 		if err != nil {
 			return err
diff --git a/blockchain/upgrade.go b/blockchain/upgrade.go
index b31f8951..cbc1d070 100644
--- a/blockchain/upgrade.go
+++ b/blockchain/upgrade.go
@@ -1804,6 +1804,25 @@ func migrateSpendJournalVersion1To2(ctx context.Context, db database.DB) error {
 	return nil
 }
 
+// initializeTreasuryBuckets creates the buckets that house the treasury account
+// and spend information as needed.
+func initializeTreasuryBuckets(db database.DB) error {
+	// Hardcoded key names so updates do not affect old upgrades.
+	treasuryBucketName := []byte("treasury")
+	treasuryTSpendBucketName := []byte("tspend")
+
+	// Create the new treasury buckets as needed.
+	return db.Update(func(dbTx database.Tx) error {
+		meta := dbTx.Metadata()
+		_, err := meta.CreateBucketIfNotExists(treasuryBucketName)
+		if err != nil {
+			return err
+		}
+		_, err = meta.CreateBucketIfNotExists(treasuryTSpendBucketName)
+		return err
+	})
+}
+
 // upgradeToVersion7 upgrades a version 6 blockchain database to version 7.
 func upgradeToVersion7(ctx context.Context, db database.DB, chainParams *chaincfg.Params, dbInfo *databaseInfo) error {
 	// Hardcoded key names so updates do not affect old upgrades.
@@ -1816,6 +1835,11 @@ func upgradeToVersion7(ctx context.Context, db database.DB, chainParams *chaincf
 	log.Info("Upgrading database to version 7...")
 	start := time.Now()
 
+	// Create the new treasury buckets as needed.
+	if err := initializeTreasuryBuckets(db); err != nil {
+		return err
+	}
+
 	// Migrate the utxoset to version 2.
 	if err := migrateUtxoSetVersion1To2(ctx, db); err != nil {
 		return err
@@ -1914,5 +1938,20 @@ func upgradeDB(ctx context.Context, db database.DB, chainParams *chaincfg.Params
 		}
 	}
 
+	// Ensure the treasury buckets are created for version 7 databases.  They
+	// ordinarily will have already been created during the upgrade to version 7
+	// above, however, due to a bug in a release candidate, they might not have
+	// been, so this is a relatively simple hack to ensure anyone in that
+	// intermediate state is upgraded properly without needing to go through
+	// another upgrade or redownload the chain.
+	//
+	// This can be removed once the database needs to migrate to version 8.
+	if dbInfo.version == 7 {
+		// Create the new treasury buckets as needed.
+		if err := initializeTreasuryBuckets(db); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }