Store a DB connection in Blockchain

blockchain.go

@@ -12,33 +12,33 @@ const blocksBucket = "blocks"
 // Blockchain keeps a sequence of Blocks
 type Blockchain struct {
 	tip []byte
+	db  *bolt.DB
 }
 
 // BlockchainIterator is used to iterate over blockchain blocks
 type BlockchainIterator struct {
 	currentHash []byte
+	db          *bolt.DB
 }
 
 // AddBlock saves provided data as a block in the blockchain
 func (bc *Blockchain) AddBlock(data string) {
 	var lastHash []byte
 
-	db, err := bolt.Open(dbFile, 0600, nil)
-	if err != nil {
-		log.Panic(err)
-	}
-	defer db.Close()
-
-	err = db.View(func(tx *bolt.Tx) error {
+	err := bc.db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(blocksBucket))
 		lastHash = b.Get([]byte("l"))
 
 		return nil
 	})
 
+	if err != nil {
+		log.Panic(err)
+	}
+
 	newBlock := NewBlock(data, lastHash)
 
-	err = db.Update(func(tx *bolt.Tx) error {
+	err = bc.db.Update(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(blocksBucket))
 		err := b.Put(newBlock.Hash, newBlock.Serialize())
 		if err != nil {
@@ -58,7 +58,7 @@ func (bc *Blockchain) AddBlock(data string) {
 
 // Iterator ...
 func (bc *Blockchain) Iterator() *BlockchainIterator {
-	bci := &BlockchainIterator{bc.tip}
+	bci := &BlockchainIterator{bc.tip, bc.db}
 
 	return bci
 }
@@ -66,13 +66,8 @@ func (bc *Blockchain) Iterator() *BlockchainIterator {
 // Next returns next block starting from the tip
 func (i *BlockchainIterator) Next() *Block {
 	var block *Block
-	db, err := bolt.Open(dbFile, 0600, nil)
-	if err != nil {
-		log.Panic(err)
-	}
-	defer db.Close()
 
-	err = db.View(func(tx *bolt.Tx) error {
+	err := i.db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(blocksBucket))
 		encodedBlock := b.Get(i.currentHash)
 		block = DeserializeBlock(encodedBlock)
@@ -80,6 +75,10 @@ func (i *BlockchainIterator) Next() *Block {
 		return nil
 	})
 
+	if err != nil {
+		log.Panic(err)
+	}
+
 	i.currentHash = block.PrevBlockHash
 
 	return block
@@ -87,12 +86,11 @@ func (i *BlockchainIterator) Next() *Block {
 
 // NewBlockchain creates a new Blockchain with genesis Block
 func NewBlockchain() *Blockchain {
-	bc := Blockchain{}
+	var tip []byte
 	db, err := bolt.Open(dbFile, 0600, nil)
 	if err != nil {
 		log.Panic(err)
 	}
-	defer db.Close()
 
 	err = db.Update(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(blocksBucket))
@@ -114,17 +112,19 @@ func NewBlockchain() *Blockchain {
 			if err != nil {
 				log.Panic(err)
 			}
-			bc.tip = genesis.Hash
+			tip = genesis.Hash
 		} else {
-			tip := b.Get([]byte("l"))
-			bc.tip = tip
+			tip = b.Get([]byte("l"))
 		}
 
 		return nil
 	})
+
 	if err != nil {
 		log.Panic(err)
 	}
 
+	bc := Blockchain{tip, db}
+
 	return &bc
 }

main.go

@@ -7,9 +7,10 @@ import (
 
 func main() {
 	bc := NewBlockchain()
+	defer bc.db.Close()
 
-	bc.AddBlock("Send 1 BTC to Ivan")
-	bc.AddBlock("Send 2 more BTC to Ivan")
+	// bc.AddBlock("Send 1 BTC to Ivan")
+	// bc.AddBlock("Send 2 more BTC to Ivan")
 
 	bci := bc.Iterator()