Convert types to take in Database interface instead of *badger.DB

We're currently taking in a concrete *badger.DB instance into most
of the types. We're soon about to introduce a garbage collecting
goroutine that would be best tested by mocking the database. Convert
the types to take in the interface instead.

7 files changed, 30 insertions, 21 deletions

diff --git a/internal/gitaly/storage/storagemgr/database.go b/internal/gitaly/storage/storagemgr/database.go
index a87e15729..14d98ff3a 100644
--- a/internal/gitaly/storage/storagemgr/database.go
+++ b/internal/gitaly/storage/storagemgr/database.go
@@ -8,14 +8,19 @@ import (
 )
 
 // OpenDatabase opens a new database handle to a database at the given path.
-func OpenDatabase(logger log.Logger, databasePath string) (*badger.DB, error) {
+func OpenDatabase(logger log.Logger, databasePath string) (Database, error) {
 	dbOptions := badger.DefaultOptions(databasePath)
 	// Enable SyncWrites to ensure all writes are persisted to disk before considering
 	// them committed.
 	dbOptions.SyncWrites = true
 	dbOptions.Logger = badgerLogger{logger}
 
-	return badger.Open(dbOptions)
+	db, err := badger.Open(dbOptions)
+	if err != nil {
+		return nil, fmt.Errorf("open: %w", err)
+	}
+
+	return newDatabaseAdapter(db), nil
 }
 
 type badgerLogger struct {
@@ -67,8 +72,7 @@ func (db databaseAdapter) GetSequence(key []byte, bandwidth uint64) (Sequence, e
 	return db.DB.GetSequence(key, bandwidth)
 }
 
-// Database is the Badger.DB interface used by TransactionManager. Refer to Badger's documentation
-// for details.
+// Database is the Badger.DB interface. Refer to Badger's documentation for details.
 type Database interface {
 	NewWriteBatch() WriteBatch
 	View(func(DatabaseTransaction) error) error
@@ -77,8 +81,7 @@ type Database interface {
 	Close() error
 }
 
-// WriteBatch is the interface of Badger.WriteBatch used by TransactionManager. Refer to Badger's
-// documentation for details.
+// WriteBatch is the interface of Badger.WriteBatch. Refer to Badger's documentation for details.
 type WriteBatch interface {
 	Set([]byte, []byte) error
 	Flush() error
diff --git a/internal/gitaly/storage/storagemgr/partition_assigner.go b/internal/gitaly/storage/storagemgr/partition_assigner.go
index 418e70b4b..5bfa45d1d 100644
--- a/internal/gitaly/storage/storagemgr/partition_assigner.go
+++ b/internal/gitaly/storage/storagemgr/partition_assigner.go
@@ -53,9 +53,9 @@ func (id partitionID) String() string {
 }
 
 // partitionAssignmentTable records which partitions repositories are assigned into.
-type partitionAssignmentTable struct{ db *badger.DB }
+type partitionAssignmentTable struct{ db Database }
 
-func newPartitionAssignmentTable(db *badger.DB) *partitionAssignmentTable {
+func newPartitionAssignmentTable(db Database) *partitionAssignmentTable {
 	return &partitionAssignmentTable{db: db}
 }
 
@@ -65,7 +65,7 @@ func (pt *partitionAssignmentTable) key(relativePath string) []byte {
 
 func (pt *partitionAssignmentTable) getPartitionID(relativePath string) (partitionID, error) {
 	var id partitionID
-	if err := pt.db.View(func(txn *badger.Txn) error {
+	if err := pt.db.View(func(txn DatabaseTransaction) error {
 		item, err := txn.Get(pt.key(relativePath))
 		if err != nil {
 			if errors.Is(err, badger.ErrKeyNotFound) {
@@ -107,7 +107,7 @@ type partitionAssigner struct {
 	// channel closing signals the lock being released.
 	repositoryLocks map[string]chan struct{}
 	// idSequence is the sequence used to mint partition IDs.
-	idSequence *badger.Sequence
+	idSequence Sequence
 	// partitionAssignmentTable contains the partition assignment records.
 	partitionAssignmentTable *partitionAssignmentTable
 	// storagePath is the path to the root directory of the storage the relative
@@ -117,7 +117,7 @@ type partitionAssigner struct {
 
 // newPartitionAssigner returns a new partitionAssigner. Close must be called on the
 // returned instance to release acquired resources.
-func newPartitionAssigner(db *badger.DB, storagePath string) (*partitionAssigner, error) {
+func newPartitionAssigner(db Database, storagePath string) (*partitionAssigner, error) {
 	seq, err := db.GetSequence([]byte("partition_id_seq"), 100)
 	if err != nil {
 		return nil, fmt.Errorf("get sequence: %w", err)
diff --git a/internal/gitaly/storage/storagemgr/partition_assigner_test.go b/internal/gitaly/storage/storagemgr/partition_assigner_test.go
index 9ccf67d98..b9239b4f0 100644
--- a/internal/gitaly/storage/storagemgr/partition_assigner_test.go
+++ b/internal/gitaly/storage/storagemgr/partition_assigner_test.go
@@ -20,11 +20,11 @@ import (
 
 type partitionAssignments map[string]partitionID
 
-func getPartitionAssignments(tb testing.TB, db *badger.DB) partitionAssignments {
+func getPartitionAssignments(tb testing.TB, db Database) partitionAssignments {
 	tb.Helper()
 
 	state := partitionAssignments{}
-	require.NoError(tb, db.View(func(txn *badger.Txn) error {
+	require.NoError(tb, db.View(func(txn DatabaseTransaction) error {
 		it := txn.NewIterator(badger.IteratorOptions{
 			Prefix: []byte(prefixPartitionAssignment),
 		})
diff --git a/internal/gitaly/storage/storagemgr/partition_manager.go b/internal/gitaly/storage/storagemgr/partition_manager.go
index 4349d5a8d..dc72f3074 100644
--- a/internal/gitaly/storage/storagemgr/partition_manager.go
+++ b/internal/gitaly/storage/storagemgr/partition_manager.go
@@ -73,7 +73,7 @@ type storageManager struct {
 	// no new transactions are allowed to begin.
 	closed bool
 	// db is the handle to the key-value store used for storing the storage's database state.
-	database *badger.DB
+	database Database
 	// partitionAssigner manages partition assignments of repositories.
 	partitionAssigner *partitionAssigner
 	// partitions contains all the active partitions. Each repository can have up to one partition.
@@ -192,6 +192,12 @@ func (ptn *partition) isClosing() bool {
 	}
 }
 
+// DatabaseOpener abstracts out the database opening for testing.
+type DatabaseOpener interface {
+	// OpenDatabase opens a database at a given path.
+	OpenDatabase(logger log.Logger, path string) (Database, error)
+}
+
 // NewPartitionManager returns a new PartitionManager.
 func NewPartitionManager(
 	configuredStorages []config.Storage,
diff --git a/internal/gitaly/storage/storagemgr/testhelper_test.go b/internal/gitaly/storage/storagemgr/testhelper_test.go
index 6fd0869d4..69e4a59cb 100644
--- a/internal/gitaly/storage/storagemgr/testhelper_test.go
+++ b/internal/gitaly/storage/storagemgr/testhelper_test.go
@@ -147,7 +147,7 @@ type DatabaseState map[string]proto.Message
 
 // RequireDatabase asserts the actual database state matches the expected database state. The actual values in the
 // database are unmarshaled to the same type the values have in the expected database state.
-func RequireDatabase(tb testing.TB, ctx context.Context, database *badger.DB, expectedState DatabaseState) {
+func RequireDatabase(tb testing.TB, ctx context.Context, database Database, expectedState DatabaseState) {
 	tb.Helper()
 
 	if expectedState == nil {
@@ -156,7 +156,7 @@ func RequireDatabase(tb testing.TB, ctx context.Context, database *badger.DB, ex
 
 	actualState := DatabaseState{}
 	unexpectedKeys := []string{}
-	require.NoError(tb, database.View(func(txn *badger.Txn) error {
+	require.NoError(tb, database.View(func(txn DatabaseTransaction) error {
 		iterator := txn.NewIterator(badger.DefaultIteratorOptions)
 		defer iterator.Close()
 
diff --git a/internal/gitaly/storage/storagemgr/transaction_manager.go b/internal/gitaly/storage/storagemgr/transaction_manager.go
index 0001c4d1a..7060eb9cd 100644
--- a/internal/gitaly/storage/storagemgr/transaction_manager.go
+++ b/internal/gitaly/storage/storagemgr/transaction_manager.go
@@ -738,7 +738,7 @@ type TransactionManager struct {
 func NewTransactionManager(
 	ptnID partitionID,
 	logger log.Logger,
-	db *badger.DB,
+	db Database,
 	storagePath,
 	relativePath,
 	stateDir,
@@ -759,7 +759,7 @@ func NewTransactionManager(
 		storagePath:          storagePath,
 		relativePath:         relativePath,
 		partitionID:          ptnID,
-		db:                   newDatabaseAdapter(db),
+		db:                   db,
 		admissionQueue:       make(chan *Transaction),
 		initialized:          make(chan struct{}),
 		snapshotLocks:        make(map[LSN]*snapshotLock),
diff --git a/internal/gitaly/storage/storagemgr/transaction_manager_hook_test.go b/internal/gitaly/storage/storagemgr/transaction_manager_hook_test.go
index e9babc31f..95fceadd7 100644
--- a/internal/gitaly/storage/storagemgr/transaction_manager_hook_test.go
+++ b/internal/gitaly/storage/storagemgr/transaction_manager_hook_test.go
@@ -19,7 +19,7 @@ type hookContext struct {
 	// closeManager calls the calls Close on the TransactionManager.
 	closeManager func()
 	// database provides access to the database for the hook handler.
-	database *badger.DB
+	database Database
 	tb       testing.TB
 }
 
@@ -42,7 +42,7 @@ type hooks struct {
 }
 
 // installHooks installs the configured hooks into the transactionManager.
-func installHooks(tb testing.TB, transactionManager *TransactionManager, database *badger.DB, hooks hooks) {
+func installHooks(tb testing.TB, transactionManager *TransactionManager, database Database, hooks hooks) {
 	hookContext := hookContext{closeManager: transactionManager.close, database: database, tb: &testingHook{TB: tb}}
 
 	transactionManager.close = func() {
@@ -62,7 +62,7 @@ func installHooks(tb testing.TB, transactionManager *TransactionManager, databas
 	}
 
 	transactionManager.db = databaseHook{
-		Database:    newDatabaseAdapter(database),
+		Database:    database,
 		hooks:       hooks,
 		hookContext: hookContext,
 	}