From 06fcf9252b679e292906ff187fbbbbf0f21a3ff7 Mon Sep 17 00:00:00 2001
From: akiozihao <akiozihaoceat@gmail.com>
Date: Tue, 9 Jan 2024 10:21:28 +0800
Subject: [PATCH] check index type

---
 db.go      |  6 +++---
 db_test.go |  9 +++++++++
 errors.go  | 17 +++++++++--------
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/db.go b/db.go
index 641b5acf..7de57b86 100644
--- a/db.go
+++ b/db.go
@@ -18,7 +18,6 @@ import (
 	"github.com/gofrs/flock"
 	"github.com/rosedblabs/diskhash"
 	"github.com/rosedblabs/wal"
-	"go.etcd.io/bbolt"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -575,6 +574,9 @@ func (db *DB) rewriteValidRecords(walFile *wal.WAL, validRecords []*ValueLogReco
 }
 
 func (db *DB) NewIterator(options IteratorOptions) (*MergeIterator, error) {
+	if db.options.IndexType != BTree {
+		return nil, ErrDBIteratorUnsupportedType
+	}
 	db.mu.Lock()
 	defer func() {
 		if r := recover(); r != nil {
@@ -583,14 +585,12 @@ func (db *DB) NewIterator(options IteratorOptions) (*MergeIterator, error) {
 	}()
 	itrs := make([]*SingleIter, 0, db.options.PartitionNum+len(db.immuMems)+1)
 	rank := 0
-	txs := make([]*bbolt.Tx, db.options.PartitionNum)
 	index := db.index.(*BPTree)
 	for i := 0; i < db.options.PartitionNum; i++ {
 		tx, err := index.trees[i].Begin(false)
 		if err != nil {
 			return nil, err
 		}
-		txs[i] = tx
 		itr, err := NewBptreeIterator(
 			tx,
 			options,
diff --git a/db_test.go b/db_test.go
index 1fa9b362..7822e2f5 100644
--- a/db_test.go
+++ b/db_test.go
@@ -786,4 +786,13 @@ func TestDBIterator(t *testing.T) {
 	}
 	err = iter.Close()
 	assert.Nil(t, err)
+
+	// unsupported type
+	options = DefaultOptions
+	options.IndexType = Hash
+	db, err = Open(options)
+	assert.Nil(t, err)
+	itr, err := db.NewIterator(IteratorOptions{Reverse: false})
+	assert.Equal(t, ErrDBIteratorUnsupportedType, err)
+	assert.Nil(t, itr)
 }
diff --git a/errors.go b/errors.go
index cfd5e76f..bb1ca424 100644
--- a/errors.go
+++ b/errors.go
@@ -3,12 +3,13 @@ package lotusdb
 import "errors"
 
 var (
-	ErrKeyIsEmpty               = errors.New("the key is empty")
-	ErrKeyNotFound              = errors.New("key not found in database")
-	ErrDatabaseIsUsing          = errors.New("the database directory is used by another process")
-	ErrReadOnlyBatch            = errors.New("the batch is read only")
-	ErrBatchCommitted           = errors.New("the batch is committed")
-	ErrDBClosed                 = errors.New("the database is closed")
-	ErrDBDirectoryISEmpty       = errors.New("the database directory path can not be empty")
-	ErrWaitMemtableSpaceTimeOut = errors.New("wait memtable space timeout, try again later")
+	ErrKeyIsEmpty                = errors.New("the key is empty")
+	ErrKeyNotFound               = errors.New("key not found in database")
+	ErrDatabaseIsUsing           = errors.New("the database directory is used by another process")
+	ErrReadOnlyBatch             = errors.New("the batch is read only")
+	ErrBatchCommitted            = errors.New("the batch is committed")
+	ErrDBClosed                  = errors.New("the database is closed")
+	ErrDBDirectoryISEmpty        = errors.New("the database directory path can not be empty")
+	ErrWaitMemtableSpaceTimeOut  = errors.New("wait memtable space timeout, try again later")
+	ErrDBIteratorUnsupportedType = errors.New("iterator only support BTree index")
 )