From b54f8a1b972858ff1d01f1a5ed2522ae08dce089 Mon Sep 17 00:00:00 2001
From: Andras Banki-Horvath <bhandras@gmail.com>
Date: Fri, 8 May 2020 17:50:13 +0200
Subject: [PATCH] kvdb+channeldb: create ExtendedBackend interface to allow
 rich tx logic

This commit adds the ExtendedBackend interface which is an extension to
the walletdb.DB interface. This paves the way to using etcd.db.View and
etcd.db.Update in the global View and Update functions without much code
rewrite.
---
 channeldb/db.go             | 33 +++++++++++++++++++++++++++++++++
 channeldb/kvdb/interface.go | 35 +++++++++++++++++++++++++++++++----
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/channeldb/db.go b/channeldb/db.go
index 9005524e..02aeb803 100644
--- a/channeldb/db.go
+++ b/channeldb/db.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/btcsuite/btcwallet/walletdb"
 	"github.com/go-errors/errors"
 	"github.com/lightningnetwork/lnd/channeldb/kvdb"
 	"github.com/lightningnetwork/lnd/channeldb/migration12"
@@ -147,12 +148,44 @@ var (
 // schedules, and reputation data.
 type DB struct {
 	kvdb.Backend
+
 	dbPath string
 	graph  *ChannelGraph
 	clock  clock.Clock
 	dryRun bool
 }
 
+// Update is a wrapper around walletdb.Update which calls into the extended
+// backend when available. This call is needed to be able to cast DB to
+// ExtendedBackend.
+func (db *DB) Update(f func(tx walletdb.ReadWriteTx) error) error {
+	if v, ok := db.Backend.(kvdb.ExtendedBackend); ok {
+		return v.Update(f)
+	}
+	return walletdb.Update(db, f)
+}
+
+// View is a wrapper around walletdb.View which calls into the extended
+// backend when available. This call is needed to be able to cast DB to
+// ExtendedBackend.
+func (db *DB) View(f func(tx walletdb.ReadTx) error) error {
+	if v, ok := db.Backend.(kvdb.ExtendedBackend); ok {
+		return v.View(f)
+	}
+
+	return walletdb.View(db, f)
+}
+
+// PrintStats calls into the extended backend if available. This call is needed
+// to be able to cast DB to ExtendedBackend.
+func (db *DB) PrintStats() string {
+	if v, ok := db.Backend.(kvdb.ExtendedBackend); ok {
+		return v.PrintStats()
+	}
+
+	return "unimplemented"
+}
+
 // Open opens or creates channeldb. Any necessary schemas migrations due
 // to updates will take place as necessary.
 // TODO(bhandras): deprecate this function.
diff --git a/channeldb/kvdb/interface.go b/channeldb/kvdb/interface.go
index 227cb006..46aded8c 100644
--- a/channeldb/kvdb/interface.go
+++ b/channeldb/kvdb/interface.go
@@ -11,13 +11,24 @@ import (
 // transaction is rolled back. If the rollback fails, the original error
 // returned by f is still returned. If the commit fails, the commit error is
 // returned.
-var Update = walletdb.Update
+func Update(db Backend, f func(tx RwTx) error) error {
+	if extendedDB, ok := db.(ExtendedBackend); ok {
+		return extendedDB.Update(f)
+	}
+	return walletdb.Update(db, f)
+}
 
 // View opens a database read transaction and executes the function f with the
 // transaction passed as a parameter. After f exits, the transaction is rolled
 // back. If f errors, its error is returned, not a rollback error (if any
 // occur).
-var View = walletdb.View
+func View(db Backend, f func(tx ReadTx) error) error {
+	if extendedDB, ok := db.(ExtendedBackend); ok {
+		return extendedDB.View(f)
+	}
+
+	return walletdb.View(db, f)
+}
 
 // Batch is identical to the Update call, but it attempts to combine several
 // individual Update transactions into a single write database transaction on
@@ -36,11 +47,27 @@ var Create = walletdb.Create
 // through read or read+write transactions.
 type Backend = walletdb.DB
 
-// BackendWithStats is and interface to debug/uncover database access patterns.
-type BackendWithStats interface {
+// ExtendedBackend is and interface that supports View and Update and also able
+// to collect database access patterns.
+type ExtendedBackend interface {
 	Backend
 
+	// PrintStats returns all collected stats pretty printed into a string.
 	PrintStats() string
+
+	// View opens a database read transaction and executes the function f with
+	// the transaction passed as a parameter. After f exits, the transaction is
+	// rolled back. If f errors, its error is returned, not a rollback error
+	// (if any occur).
+	View(f func(tx walletdb.ReadTx) error) error
+
+	// Update opens a database read/write transaction and executes the function
+	// f with the transaction passed as a parameter. After f exits, if f did not
+	// error, the transaction is committed. Otherwise, if f did error, the
+	// transaction is rolled back. If the rollback fails, the original error
+	// returned by f is still returned. If the commit fails, the commit error is
+	// returned.
+	Update(f func(tx walletdb.ReadWriteTx) error) error
 }
 
 // Open opens an existing database for the specified type. The arguments are