From efeb7ff7e65badf16db85dea64dde242afddc791 Mon Sep 17 00:00:00 2001
From: Roy Crihfield <roy@manteia.ltd>
Date: Fri, 10 Mar 2023 15:54:39 +0800
Subject: [PATCH] clean up

---
 pkg/database.go       | 41 +++++++++++++++++++++++++++++++++--------
 pkg/util.go           | 37 -------------------------------------
 pkg/validator.go      |  3 ++-
 pkg/validator_test.go |  1 +
 4 files changed, 36 insertions(+), 46 deletions(-)

diff --git a/pkg/database.go b/pkg/database.go
index 3f34ccf..be2b756 100644
--- a/pkg/database.go
+++ b/pkg/database.go
@@ -18,6 +18,8 @@ package validator
 
 import (
 	"fmt"
+	"os"
+	"strconv"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/spf13/viper"
@@ -32,13 +34,6 @@ const (
 	DATABASE_PASSWORD = "DATABASE_PASSWORD"
 )
 
-// NewDB returns a new sqlx.DB from config/cli/env variables
-func NewDB() (*sqlx.DB, error) {
-	c := Config{}
-	c.Init()
-	return sqlx.Connect("postgres", c.ConnString())
-}
-
 type Config struct {
 	Hostname string
 	Name     string
@@ -47,6 +42,13 @@ type Config struct {
 	Port     int
 }
 
+// NewDB returns a new sqlx.DB from config/cli/env variables
+func NewDB() (*sqlx.DB, error) {
+	c := Config{}
+	LoadViper(&c)
+	return sqlx.Connect("postgres", c.ConnString())
+}
+
 func (c *Config) ConnString() string {
 	if len(c.User) > 0 && len(c.Password) > 0 {
 		return fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=disable",
@@ -59,7 +61,30 @@ func (c *Config) ConnString() string {
 	return fmt.Sprintf("postgresql://%s:%d/%s?sslmode=disable", c.Hostname, c.Port, c.Name)
 }
 
-func (c *Config) Init() {
+func LoadEnv(c *Config) error {
+	if val := os.Getenv(DATABASE_NAME); val != "" {
+		c.Name = val
+	}
+	if val := os.Getenv(DATABASE_HOSTNAME); val != "" {
+		c.Hostname = val
+	}
+	if val := os.Getenv(DATABASE_PORT); val != "" {
+		port, err := strconv.Atoi(val)
+		if err != nil {
+			return err
+		}
+		c.Port = port
+	}
+	if val := os.Getenv(DATABASE_USER); val != "" {
+		c.User = val
+	}
+	if val := os.Getenv(DATABASE_PASSWORD); val != "" {
+		c.Password = val
+	}
+	return nil
+}
+
+func LoadViper(c *Config) {
 	viper.BindEnv("database.name", DATABASE_NAME)
 	viper.BindEnv("database.hostname", DATABASE_HOSTNAME)
 	viper.BindEnv("database.port", DATABASE_PORT)
diff --git a/pkg/util.go b/pkg/util.go
index 9cfd46e..5a2ca8e 100644
--- a/pkg/util.go
+++ b/pkg/util.go
@@ -20,41 +20,10 @@ import (
 	"context"
 
 	"github.com/ipfs/go-blockservice"
-	"github.com/ipfs/go-cid"
-	blockstore "github.com/ipfs/go-ipfs-blockstore"
-	dshelp "github.com/ipfs/go-ipfs-ds-help"
 	"github.com/ipfs/kubo/core"
 	"github.com/ipfs/kubo/repo/fsrepo"
-	"github.com/jmoiron/sqlx"
 )
 
-// PublishRaw derives a cid from raw bytes and provided codec and multihash type, and writes it to the db tx
-func PublishRaw(tx *sqlx.Tx, codec, mh uint64, raw []byte, blockNumber uint64) (string, error) {
-	c, err := RawdataToCid(codec, raw, mh)
-	if err != nil {
-		return "", err
-	}
-	dbKey := dshelp.MultihashToDsKey(c.Hash())
-	prefixedKey := blockstore.BlockPrefix.String() + dbKey.String()
-	_, err = tx.Exec(`INSERT INTO public.blocks (key, data, block_number) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING`, prefixedKey, raw, blockNumber)
-	return c.String(), err
-}
-
-// RawdataToCid takes the desired codec, multihash type, and a slice of bytes
-// and returns the proper cid of the object.
-func RawdataToCid(codec uint64, rawdata []byte, multiHash uint64) (cid.Cid, error) {
-	c, err := cid.Prefix{
-		Codec:    codec,
-		Version:  1,
-		MhType:   multiHash,
-		MhLength: -1,
-	}.Sum(rawdata)
-	if err != nil {
-		return cid.Cid{}, err
-	}
-	return c, nil
-}
-
 // InitIPFSBlockService is used to configure and return a BlockService using an ipfs repo path (e.g. ~/.ipfs)
 func InitIPFSBlockService(ipfsPath string) (blockservice.BlockService, error) {
 	r, openErr := fsrepo.Open(ipfsPath)
@@ -72,9 +41,3 @@ func InitIPFSBlockService(ipfsPath string) (blockservice.BlockService, error) {
 	}
 	return ipfsNode.Blocks, nil
 }
-
-// ResetTestDB drops all rows in the test db public.blocks table
-func ResetTestDB(db *sqlx.DB) error {
-	_, err := db.Exec("DELETE FROM public.blocks")
-	return err
-}
diff --git a/pkg/validator.go b/pkg/validator.go
index 63b282b..fed96c1 100644
--- a/pkg/validator.go
+++ b/pkg/validator.go
@@ -205,7 +205,8 @@ func (v *Validator) iterate(it trie.NodeIterator, storage bool) error {
 	return it.Error()
 }
 
-// Traverses each iterator in a separate goroutine, dumping to a recovery file on failure. Handles signals
+// Traverses each iterator in a separate goroutine.
+// Dumps to a recovery file on failure or interrupt.
 func iterateTracked(tree state.Trie, recoveryFile string, iterCount uint, fn func(trie.NodeIterator) error) error {
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	tracker := tracker.New(recoveryFile, iterCount)
diff --git a/pkg/validator_test.go b/pkg/validator_test.go
index d530dc4..731be27 100644
--- a/pkg/validator_test.go
+++ b/pkg/validator_test.go
@@ -224,6 +224,7 @@ var _ = Describe("PG-IPFS Validator", func() {
 	AfterEach(func() {
 		os.RemoveAll(tmp)
 		v.Close()
+		db.Close()
 	})
 	Describe("ValidateTrie", func() {
 		AfterEach(func() {