From e80e21d1a8102c45223ad7a0f6fe6d41176ca49e Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Thu, 21 May 2020 15:36:26 -0700
Subject: [PATCH] channeldb/migration: add generic CreateTLB migration

This commit creates a generic migration for creating top-level buckets.
---
 channeldb/log.go                       |  2 +
 channeldb/migration/create_tlb.go      | 27 ++++++++++++
 channeldb/migration/create_tlb_test.go | 57 ++++++++++++++++++++++++++
 channeldb/migration/log.go             | 12 ++++++
 4 files changed, 98 insertions(+)
 create mode 100644 channeldb/migration/create_tlb.go
 create mode 100644 channeldb/migration/create_tlb_test.go
 create mode 100644 channeldb/migration/log.go

diff --git a/channeldb/log.go b/channeldb/log.go
index 7490c6bf5..f59426f0a 100644
--- a/channeldb/log.go
+++ b/channeldb/log.go
@@ -3,6 +3,7 @@ package channeldb
 import (
 	"github.com/btcsuite/btclog"
 	"github.com/lightningnetwork/lnd/build"
+	mig "github.com/lightningnetwork/lnd/channeldb/migration"
 	"github.com/lightningnetwork/lnd/channeldb/migration12"
 	"github.com/lightningnetwork/lnd/channeldb/migration13"
 	"github.com/lightningnetwork/lnd/channeldb/migration_01_to_11"
@@ -28,6 +29,7 @@ func DisableLog() {
 // using btclog.
 func UseLogger(logger btclog.Logger) {
 	log = logger
+	mig.UseLogger(logger)
 	migration_01_to_11.UseLogger(logger)
 	migration12.UseLogger(logger)
 	migration13.UseLogger(logger)
diff --git a/channeldb/migration/create_tlb.go b/channeldb/migration/create_tlb.go
new file mode 100644
index 000000000..7c31ec929
--- /dev/null
+++ b/channeldb/migration/create_tlb.go
@@ -0,0 +1,27 @@
+package migration
+
+import (
+	"fmt"
+
+	"github.com/lightningnetwork/lnd/channeldb/kvdb"
+)
+
+// CreateTLB creates a new top-level bucket with the passed bucket identifier.
+func CreateTLB(bucket []byte) func(kvdb.RwTx) error {
+	return func(tx kvdb.RwTx) error {
+		log.Infof("Creating top-level bucket: \"%s\" ...", bucket)
+
+		if tx.ReadBucket(bucket) != nil {
+			return fmt.Errorf("top-level bucket \"%s\" "+
+				"already exists", bucket)
+		}
+
+		_, err := tx.CreateTopLevelBucket(bucket)
+		if err != nil {
+			return err
+		}
+
+		log.Infof("Created top-level bucket: \"%s\"", bucket)
+		return nil
+	}
+}
diff --git a/channeldb/migration/create_tlb_test.go b/channeldb/migration/create_tlb_test.go
new file mode 100644
index 000000000..2f4221361
--- /dev/null
+++ b/channeldb/migration/create_tlb_test.go
@@ -0,0 +1,57 @@
+package migration_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/lightningnetwork/lnd/channeldb/kvdb"
+	"github.com/lightningnetwork/lnd/channeldb/migration"
+	"github.com/lightningnetwork/lnd/channeldb/migtest"
+)
+
+// TestCreateTLB asserts that a CreateTLB properly initializes a new top-level
+// bucket, and that it succeeds even if the bucket already exists. It would
+// probably be better if the latter failed, but the kvdb abstraction doesn't
+// support this.
+func TestCreateTLB(t *testing.T) {
+	newBucket := []byte("hello")
+
+	tests := []struct {
+		name            string
+		beforeMigration func(kvdb.RwTx) error
+		shouldFail      bool
+	}{
+		{
+			name: "already exists",
+			beforeMigration: func(tx kvdb.RwTx) error {
+				_, err := tx.CreateTopLevelBucket(newBucket)
+				return err
+			},
+			shouldFail: true,
+		},
+		{
+			name:            "does not exist",
+			beforeMigration: func(_ kvdb.RwTx) error { return nil },
+			shouldFail:      false,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			migtest.ApplyMigration(
+				t,
+				test.beforeMigration,
+				func(tx kvdb.RwTx) error {
+					if tx.ReadBucket(newBucket) != nil {
+						return nil
+					}
+					return fmt.Errorf("bucket \"%s\" not "+
+						"created", newBucket)
+				},
+				migration.CreateTLB(newBucket),
+				test.shouldFail,
+			)
+		})
+	}
+}
diff --git a/channeldb/migration/log.go b/channeldb/migration/log.go
new file mode 100644
index 000000000..5085596da
--- /dev/null
+++ b/channeldb/migration/log.go
@@ -0,0 +1,12 @@
+package migration
+
+import "github.com/btcsuite/btclog"
+
+// log is a logger that is initialized as disabled.  This means the package will
+// not perform any logging by default until a logger is set.
+var log = btclog.Disabled
+
+// UseLogger uses a specified Logger to output package logging info.
+func UseLogger(logger btclog.Logger) {
+	log = logger
+}