Fix insert statements and upgrade dependencies (#31)

Upgrade vulcanize-geth to v4

.gitignore

@@ -1,2 +1,2 @@
 .idea/
-eth-pg-ipfs-state-snapshot
+ipld-eth-state-snapshot

Makefile

@@ -6,7 +6,7 @@ mockgen_cmd=mockgen
 mocks: mocks/snapshot/publisher.go
 
 mocks/snapshot/publisher.go: pkg/types/publisher.go
-	$(mockgen_cmd) -package snapshot_mock -destination $@ -source $< Publisher
+	$(mockgen_cmd) -package snapshot_mock -destination $@ -source $< Publisher Tx
 
 clean:
 	rm -f mocks/snapshot/publisher.go

README.md

@@ -1,26 +1,35 @@
-# eth-pg-ipfs-state-snapshot
+# ipld-eth-state-snapshot
 
 > Tool for extracting the entire Ethereum state at a particular block height from leveldb into Postgres-backed IPFS
 
-[![Go Report Card](https://goreportcard.com/badge/github.com/vulcanize/eth-pg-ipfs-state-snapshot)](https://goreportcard.com/report/github.com/vulcanize/eth-pg-ipfs-state-snapshot)
+[![Go Report Card](https://goreportcard.com/badge/github.com/vulcanize/ipld-eth-state-snapshot)](https://goreportcard.com/report/github.com/vulcanize/ipld-eth-state-snapshot)
 
 ## Usage 
 
-./eth-pg-ipfs-state-snapshot stateSnapshot --config={path to toml config file}
+./ipld-eth-state-snapshot stateSnapshot --config={path to toml config file}
+
+### Config
 
 Config format:
 
 ```toml
-[database]
-    name     = "vulcanize_public"
-    hostname = "localhost"
-    port     = 5432
-    user     = "postgres"
+[snapshot]
+    mode = "file" # indicates output mode ("postgres" or "file")
+    workers = 4 # degree of concurrency, the state trie is subdivided into sectiosn that are traversed and processed concurrently
+    blockHeight = -1 # blockheight to perform the snapshot at (-1 indicates to use the latest blockheight found in leveldb)
+    recoveryFile = "recovery_file" # specifies a file to output recovery information on error or premature closure
 
 [leveldb]
-   path = "/Users/user/Library/Ethereum/geth/chaindata"
-   ancient = "/Users/user/Library/Ethereum/geth/chaindata/ancient"
+    path = "/Users/user/Library/Ethereum/geth/chaindata" # path to geth leveldb
+    ancient = "/Users/user/Library/Ethereum/geth/chaindata/ancient" # path to geth ancient database
 
-[snapshot]
-   blockHeight = 0
+[database]
+    name     = "vulcanize_public" # postgres database name
+    hostname = "localhost" # postgres host
+    port     = 5432 # postgres port
+    user     = "postgres" # postgres user
+    password = "" # postgres password
+
+[file]
+    outputDir = "output_dir/" # when operating in 'file' output mode, this is the directory the files are written to
 ```

cmd/root.go

@@ -33,7 +33,7 @@ var (
 )
 
 var rootCmd = &cobra.Command{
-	Use:              "eth-pg-ipfs-state-snapshot",
+	Use:              "ipld-eth-state-snapshot",
 	PersistentPreRun: initFuncs,
 }
 

cmd/stateSnapshot.go

@@ -16,11 +16,13 @@
 package cmd
 
 import (
+	"fmt"
+
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 
-	"github.com/vulcanize/eth-pg-ipfs-state-snapshot/pkg/snapshot"
+	"github.com/vulcanize/ipld-eth-state-snapshot/pkg/snapshot"
 )
 
 // stateSnapshotCmd represents the stateSnapshot command
@@ -29,7 +31,7 @@ var stateSnapshotCmd = &cobra.Command{
 	Short: "Extract the entire Ethereum state from leveldb and publish into PG-IPFS",
 	Long: `Usage
 
-./eth-pg-ipfs-state-snapshot stateSnapshot --config={path to toml config file}`,
+./ipld-eth-state-snapshot stateSnapshot --config={path to toml config file}`,
 	Run: func(cmd *cobra.Command, args []string) {
 		subCommand = cmd.CalledAs()
 		logWithCommand = *logrus.WithField("SubCommand", subCommand)
@@ -38,22 +40,36 @@ var stateSnapshotCmd = &cobra.Command{
 }
 
 func stateSnapshot() {
-	snapConfig := snapshot.NewConfig()
-	pgDB, err := snapshot.NewPostgresDB(snapConfig.DB)
+	modeStr := viper.GetString("snapshot.mode")
+	mode := snapshot.SnapshotMode(modeStr)
+	config, err := snapshot.NewConfig(mode)
 	if err != nil {
-		logWithCommand.Fatal(err)
+		logWithCommand.Fatal("unable to initialize config: %v", err)
 	}
-	edb, err := snapshot.NewLevelDB(snapConfig.Eth)
-	if err != nil {
-		logWithCommand.Fatal(err)
-	}
-
-	snapshotService, err := snapshot.NewSnapshotService(edb, snapshot.NewPublisher(pgDB))
+	logWithCommand.Infof("opening levelDB and ancient data at %s and %s",
+		config.Eth.LevelDBPath, config.Eth.AncientDBPath)
+	edb, err := snapshot.NewLevelDB(config.Eth)
 	if err != nil {
 		logWithCommand.Fatal(err)
 	}
 	height := viper.GetInt64("snapshot.blockHeight")
+	recoveryFile := viper.GetString("snapshot.recoveryFile")
+	if recoveryFile == "" {
+		recoveryFile = fmt.Sprintf("./%d_snapshot_recovery", height)
+		logWithCommand.Infof("no recovery file set, creating default: %s", recoveryFile)
+	}
+
+	pub, err := snapshot.NewPublisher(mode, config)
+	if err != nil {
+		logWithCommand.Fatal(err)
+	}
+
+	snapshotService, err := snapshot.NewSnapshotService(edb, pub, recoveryFile)
+	if err != nil {
+		logWithCommand.Fatal(err)
+	}
 	workers := viper.GetUint("snapshot.workers")
+
 	if height < 0 {
 		if err := snapshotService.CreateLatestSnapshot(workers); err != nil {
 			logWithCommand.Fatal(err)
@@ -73,10 +89,16 @@ func init() {
 	stateSnapshotCmd.PersistentFlags().String("leveldb-path", "", "path to primary datastore")
 	stateSnapshotCmd.PersistentFlags().String("ancient-path", "", "path to ancient datastore")
 	stateSnapshotCmd.PersistentFlags().String("block-height", "", "blockheight to extract state at")
-	stateSnapshotCmd.PersistentFlags().Int("workers", 0, "number of concurrent workers to use")
+	stateSnapshotCmd.PersistentFlags().Int("workers", 1, "number of concurrent workers to use")
+	stateSnapshotCmd.PersistentFlags().String("recovery-file", "", "file to recover from a previous iteration")
+	stateSnapshotCmd.PersistentFlags().String("snapshot-mode", "postgres", "output mode for snapshot ('file' or 'postgres')")
+	stateSnapshotCmd.PersistentFlags().String("output-dir", "", "directory for writing ouput to while operating in 'file' mode")
 
 	viper.BindPFlag("leveldb.path", stateSnapshotCmd.PersistentFlags().Lookup("leveldb-path"))
 	viper.BindPFlag("leveldb.ancient", stateSnapshotCmd.PersistentFlags().Lookup("ancient-path"))
 	viper.BindPFlag("snapshot.blockHeight", stateSnapshotCmd.PersistentFlags().Lookup("block-height"))
 	viper.BindPFlag("snapshot.workers", stateSnapshotCmd.PersistentFlags().Lookup("workers"))
+	viper.BindPFlag("snapshot.recoveryFile", stateSnapshotCmd.PersistentFlags().Lookup("recovery-file"))
+	viper.BindPFlag("snapshot.mode", stateSnapshotCmd.PersistentFlags().Lookup("snapshot-mode"))
+	viper.BindPFlag("file.outputDir", stateSnapshotCmd.PersistentFlags().Lookup("output-dir"))
 }

environments/example.toml

@@ -9,4 +9,10 @@
     ancient = "/Users/user/Library/Ethereum/geth/chaindata/ancient"
 
 [snapshot]
-    blockHeight = -1
+    mode = "file"
+    workers = 4
+    blockHeight = -1
+    recoveryFile = "recovery_file"
+
+[file]
+    outputDir = "output_dir/"

fixture/chaindata.go

@@ -0,0 +1,27 @@
+package fixture
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+)
+
+// TODO: embed some mainnet data
+// import "embed"
+//_go:embed mainnet_data.tar.gz
+
+var (
+	ChaindataPath, AncientdataPath string
+)
+
+func init() {
+	_, path, _, _ := runtime.Caller(0)
+	wd := filepath.Dir(path)
+
+	ChaindataPath = filepath.Join(wd, "..", "fixture", "chaindata")
+	AncientdataPath = filepath.Join(ChaindataPath, "ancient")
+
+	if _, err := os.Stat(ChaindataPath); err != nil {
+		panic("must populate chaindata at " + ChaindataPath)
+	}
+}

fixture/chaindata/.gitignore

@@ -0,0 +1,6 @@
+*.log
+CURRENT*
+LOCK
+LOG
+MANIFEST-*
+ancient/FLOCK

fixture/node_paths.go

@@ -0,0 +1,359 @@
+package fixture
+
+var Block1_StateNodePaths = [][]byte{
+	[]byte{},
+	[]byte{0},
+	[]byte{0, 0},
+	[]byte{0, 2},
+	[]byte{0, 2, 1},
+	[]byte{0, 2, 8},
+	[]byte{0, 2, 12},
+	[]byte{0, 3},
+	[]byte{0, 4},
+	[]byte{0, 6},
+	[]byte{0, 6, 3},
+	[]byte{0, 6, 13},
+	[]byte{0, 7},
+	[]byte{0, 8},
+	[]byte{0, 8, 7},
+	[]byte{0, 8, 11},
+	[]byte{0, 9},
+	[]byte{0, 9, 9},
+	[]byte{0, 9, 10},
+	[]byte{0, 12},
+	[]byte{0, 13},
+	[]byte{0, 14},
+	[]byte{1},
+	[]byte{1, 2},
+	[]byte{1, 2, 5},
+	[]byte{1, 2, 7},
+	[]byte{1, 3},
+	[]byte{1, 3, 1},
+	[]byte{1, 3, 11},
+	[]byte{1, 4},
+	[]byte{1, 5},
+	[]byte{1, 5, 11},
+	[]byte{1, 5, 12},
+	[]byte{1, 5, 15},
+	[]byte{1, 6},
+	[]byte{1, 8},
+	[]byte{1, 10},
+	[]byte{1, 13},
+	[]byte{1, 14},
+	[]byte{1, 14, 2},
+	[]byte{1, 14, 11},
+	[]byte{1, 15},
+	[]byte{1, 15, 9},
+	[]byte{1, 15, 15},
+	[]byte{2},
+	[]byte{2, 0},
+	[]byte{2, 0, 9},
+	[]byte{2, 0, 14},
+	[]byte{2, 1},
+	[]byte{2, 1, 1},
+	[]byte{2, 1, 3},
+	[]byte{2, 1, 14},
+	[]byte{2, 5},
+	[]byte{2, 6},
+	[]byte{2, 9},
+	[]byte{2, 9, 1},
+	[]byte{2, 9, 7},
+	[]byte{2, 11},
+	[]byte{2, 11, 7},
+	[]byte{2, 11, 13},
+	[]byte{2, 13},
+	[]byte{2, 13, 1},
+	[]byte{2, 13, 15},
+	[]byte{2, 15},
+	[]byte{3},
+	[]byte{3, 0},
+	[]byte{3, 0, 0},
+	[]byte{3, 0, 1},
+	[]byte{3, 2},
+	[]byte{3, 2, 3},
+	[]byte{3, 2, 15},
+	[]byte{3, 3},
+	[]byte{3, 4},
+	[]byte{3, 4, 2},
+	[]byte{3, 4, 4},
+	[]byte{3, 4, 5},
+	[]byte{3, 6},
+	[]byte{3, 8},
+	[]byte{3, 9},
+	[]byte{3, 10},
+	[]byte{3, 10, 2},
+	[]byte{3, 10, 8},
+	[]byte{3, 10, 12},
+	[]byte{3, 11},
+	[]byte{3, 12},
+	[]byte{3, 13},
+	[]byte{3, 14},
+	[]byte{3, 14, 4},
+	[]byte{3, 14, 9},
+	[]byte{3, 14, 14},
+	[]byte{3, 14, 14, 10},
+	[]byte{3, 14, 14, 15},
+	[]byte{4},
+	[]byte{4, 0},
+	[]byte{4, 0, 6},
+	[]byte{4, 0, 15},
+	[]byte{4, 1},
+	[]byte{4, 2},
+	[]byte{4, 2, 1},
+	[]byte{4, 2, 11},
+	[]byte{4, 3},
+	[]byte{4, 5},
+	[]byte{4, 6},
+	[]byte{4, 7},
+	[]byte{4, 8},
+	[]byte{4, 11},
+	[]byte{4, 11, 6},
+	[]byte{4, 11, 9},
+	[]byte{4, 11, 12},
+	[]byte{4, 14},
+	[]byte{5},
+	[]byte{5, 0},
+	[]byte{5, 0, 3},
+	[]byte{5, 0, 9},
+	[]byte{5, 0, 15},
+	[]byte{5, 1},
+	[]byte{5, 1, 14},
+	[]byte{5, 1, 15},
+	[]byte{5, 2},
+	[]byte{5, 2, 8},
+	[]byte{5, 2, 10},
+	[]byte{5, 3},
+	[]byte{5, 4},
+	[]byte{5, 4, 6},
+	[]byte{5, 4, 12},
+	[]byte{5, 6},
+	[]byte{5, 8},
+	[]byte{5, 8, 3},
+	[]byte{5, 8, 11},
+	[]byte{5, 10},
+	[]byte{5, 11},
+	[]byte{5, 12},
+	[]byte{5, 13},
+	[]byte{5, 15},
+	[]byte{6},
+	[]byte{6, 0},
+	[]byte{6, 2},
+	[]byte{6, 2, 3},
+	[]byte{6, 2, 9},
+	[]byte{6, 4},
+	[]byte{6, 4, 0},
+	[]byte{6, 4, 0, 0},
+	[]byte{6, 4, 0, 5},
+	[]byte{6, 5},
+	[]byte{6, 5, 4},
+	[]byte{6, 5, 10},
+	[]byte{6, 5, 12},
+	[]byte{6, 5, 13},
+	[]byte{6, 6},
+	[]byte{6, 6, 0},
+	[]byte{6, 6, 8},
+	[]byte{6, 8},
+	[]byte{6, 8, 4},
+	[]byte{6, 8, 4, 2},
+	[]byte{6, 8, 4, 9},
+	[]byte{6, 8, 9},
+	[]byte{6, 10},
+	[]byte{6, 10, 1},
+	[]byte{6, 10, 14},
+	[]byte{6, 11},
+	[]byte{6, 11, 2},
+	[]byte{6, 11, 12},
+	[]byte{6, 11, 14},
+	[]byte{6, 13},
+	[]byte{6, 13, 2},
+	[]byte{6, 13, 12},
+	[]byte{7},
+	[]byte{7, 1},
+	[]byte{7, 5},
+	[]byte{7, 7},
+	[]byte{7, 8},
+	[]byte{7, 8, 2},
+	[]byte{7, 8, 5},
+	[]byte{7, 9},
+	[]byte{7, 13},
+	[]byte{7, 13, 1},
+	[]byte{7, 13, 1, 0},
+	[]byte{7, 13, 1, 13},
+	[]byte{7, 13, 7},
+	[]byte{7, 14},
+	[]byte{7, 14, 8},
+	[]byte{7, 14, 11},
+	[]byte{8},
+	[]byte{8, 0},
+	[]byte{8, 0, 3},
+	[]byte{8, 0, 11},
+	[]byte{8, 2},
+	[]byte{8, 4},
+	[]byte{8, 8},
+	[]byte{8, 9},
+	[]byte{8, 9, 3},
+	[]byte{8, 9, 13},
+	[]byte{8, 10},
+	[]byte{8, 12},
+	[]byte{8, 12, 3},
+	[]byte{8, 12, 15},
+	[]byte{8, 13},
+	[]byte{8, 15},
+	[]byte{8, 15, 8},
+	[]byte{8, 15, 13},
+	[]byte{9},
+	[]byte{9, 0},
+	[]byte{9, 5},
+	[]byte{9, 6},
+	[]byte{9, 6, 10},
+	[]byte{9, 6, 14},
+	[]byte{9, 7},
+	[]byte{9, 9},
+	[]byte{9, 14},
+	[]byte{9, 15},
+	[]byte{9, 15, 0},
+	[]byte{9, 15, 4},
+	[]byte{9, 15, 10},
+	[]byte{10},
+	[]byte{10, 0},
+	[]byte{10, 0, 9},
+	[]byte{10, 0, 10},
+	[]byte{10, 0, 15},
+	[]byte{10, 2},
+	[]byte{10, 3},
+	[]byte{10, 6},
+	[]byte{10, 8},
+	[]byte{10, 9},
+	[]byte{10, 10},
+	[]byte{10, 10, 5},
+	[]byte{10, 10, 8},
+	[]byte{10, 13},
+	[]byte{10, 13, 0},
+	[]byte{10, 13, 13},
+	[]byte{10, 14},
+	[]byte{10, 14, 4},
+	[]byte{10, 14, 11},
+	[]byte{10, 14, 11, 8},
+	[]byte{10, 14, 11, 14},
+	[]byte{10, 15},
+	[]byte{11},
+	[]byte{11, 0},
+	[]byte{11, 0, 2},
+	[]byte{11, 0, 15},
+	[]byte{11, 1},
+	[]byte{11, 2},
+	[]byte{11, 3},
+	[]byte{11, 4},
+	[]byte{11, 5},
+	[]byte{11, 7},
+	[]byte{11, 7, 12},
+	[]byte{11, 7, 15},
+	[]byte{11, 8},
+	[]byte{11, 8, 8},
+	[]byte{11, 8, 15},
+	[]byte{11, 9},
+	[]byte{11, 11},
+	[]byte{11, 12},
+	[]byte{11, 13},
+	[]byte{11, 14},
+	[]byte{11, 14, 0},
+	[]byte{11, 14, 0, 1},
+	[]byte{11, 14, 0, 3},
+	[]byte{11, 14, 8},
+	[]byte{11, 14, 13},
+	[]byte{12},
+	[]byte{12, 0},
+	[]byte{12, 0, 0},
+	[]byte{12, 0, 1},
+	[]byte{12, 0, 1, 3},
+	[]byte{12, 0, 1, 11},
+	[]byte{12, 0, 15},
+	[]byte{12, 2},
+	[]byte{12, 2, 9},
+	[]byte{12, 2, 12},
+	[]byte{12, 4},
+	[]byte{12, 5},
+	[]byte{12, 6},
+	[]byte{12, 6, 0},
+	[]byte{12, 6, 4},
+	[]byte{12, 6, 14},
+	[]byte{12, 7},
+	[]byte{12, 7, 0},
+	[]byte{12, 7, 12},
+	[]byte{12, 7, 13},
+	[]byte{12, 9},
+	[]byte{12, 11},
+	[]byte{12, 12},
+	[]byte{13},
+	[]byte{13, 2},
+	[]byte{13, 2, 0},
+	[]byte{13, 2, 2},
+	[]byte{13, 2, 4},
+	[]byte{13, 3},
+	[]byte{13, 3, 7},
+	[]byte{13, 3, 10},
+	[]byte{13, 5},
+	[]byte{13, 8},
+	[]byte{13, 8, 1},
+	[]byte{13, 8, 15},
+	[]byte{13, 9},
+	[]byte{13, 9, 0},
+	[]byte{13, 9, 14},
+	[]byte{13, 10},
+	[]byte{13, 12},
+	[]byte{13, 12, 8},
+	[]byte{13, 12, 11},
+	[]byte{13, 13},
+	[]byte{13, 13, 7},
+	[]byte{13, 13, 12},
+	[]byte{13, 14},
+	[]byte{14},
+	[]byte{14, 0},
+	[]byte{14, 1},
+	[]byte{14, 2},
+	[]byte{14, 2, 2},
+	[]byte{14, 2, 12},
+	[]byte{14, 3},
+	[]byte{14, 4},
+	[]byte{14, 5},
+	[]byte{14, 6},
+	[]byte{14, 6, 9},
+	[]byte{14, 6, 12},
+	[]byte{14, 7},
+	[]byte{14, 7, 4},
+	[]byte{14, 7, 12},
+	[]byte{14, 8},
+	[]byte{14, 8, 3},
+	[]byte{14, 8, 12},
+	[]byte{14, 8, 12, 0},
+	[]byte{14, 8, 12, 6},
+	[]byte{14, 10},
+	[]byte{14, 10, 6},
+	[]byte{14, 10, 12},
+	[]byte{14, 11},
+	[]byte{14, 11, 8},
+	[]byte{14, 11, 13},
+	[]byte{14, 12},
+	[]byte{14, 14},
+	[]byte{14, 14, 3},
+	[]byte{14, 14, 9},
+	[]byte{15},
+	[]byte{15, 0},
+	[]byte{15, 5},
+	[]byte{15, 6},
+	[]byte{15, 9},
+	[]byte{15, 9, 0},
+	[]byte{15, 9, 2},
+	[]byte{15, 9, 3},
+	[]byte{15, 11},
+	[]byte{15, 11, 1},
+	[]byte{15, 11, 6},
+	[]byte{15, 12},
+	[]byte{15, 12, 3},
+	[]byte{15, 12, 14},
+	[]byte{15, 12, 14, 7},
+	[]byte{15, 12, 14, 13},
+	[]byte{15, 13},
+	[]byte{15, 14},
+	[]byte{15, 15},
+}

fixture/service.go

@@ -5,12 +5,11 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
-	// "github.com/jmoiron/sqlx"
 
-	snapt "github.com/vulcanize/eth-pg-ipfs-state-snapshot/pkg/types"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
 )
 
-var PublishHeader = &types.Header{
+var Block1_Header = types.Header{
 	ParentHash:  common.HexToHash("0x6341fd3daf94b748c72ced5a5b26028f2474f5f00d824504e4fa37a75767e177"),
 	UncleHash:   common.HexToHash("0x1dcc4de8dec75d7aab85b567b6ccd41ad312451b948a7413f0a142fd40d49347"),
 	Coinbase:    common.HexToAddress("0x0000000000000000000000000000000000000000"),
@@ -29,7 +28,7 @@ var PublishHeader = &types.Header{
 	BaseFee:     nil,
 }
 
-var StateNode = snapt.Node{
+var Block1_StateNode0 = snapt.Node{
 	NodeType: 0,
 	Path:     []byte{12, 0},
 	Key:      common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000000"),

go.mod

@@ -1,45 +1,20 @@
-module github.com/vulcanize/eth-pg-ipfs-state-snapshot
+module github.com/vulcanize/ipld-eth-state-snapshot
 
 go 1.15
 
 require (
-	github.com/btcsuite/btcd v0.22.0-beta // indirect
-	github.com/cespare/xxhash/v2 v2.1.2 // indirect
-	github.com/ethereum/go-ethereum v1.10.15
-	github.com/fsnotify/fsnotify v1.5.1 // indirect
-	github.com/go-kit/kit v0.10.0 // indirect
+	github.com/btcsuite/btcd/btcec/v2 v2.2.0 // indirect
+	github.com/ethereum/go-ethereum v1.10.17
 	github.com/golang/mock v1.6.0
-	github.com/google/go-cmp v0.5.6 // indirect
-	github.com/google/uuid v1.3.0 // indirect
-	github.com/gopherjs/gopherjs v0.0.0-20190430165422-3e4dfb77656c // indirect
-	github.com/ipfs/go-cid v0.1.0 // indirect
-	github.com/ipfs/go-datastore v0.5.1 // indirect
-	github.com/ipfs/go-ipfs-blockstore v1.1.2 // indirect
-	github.com/ipfs/go-log v1.0.5 // indirect
-	github.com/ipfs/go-log/v2 v2.4.0 // indirect
-	github.com/jmoiron/sqlx v1.2.0
-	github.com/kr/pretty v0.3.0 // indirect
-	github.com/multiformats/go-base32 v0.0.4 // indirect
+	github.com/ipfs/go-cid v0.1.0
+	github.com/ipfs/go-ipfs-blockstore v1.1.2
+	github.com/ipfs/go-ipfs-ds-help v1.1.0
+	github.com/jackc/pgx/v4 v4.15.0
 	github.com/multiformats/go-multihash v0.1.0
-	github.com/onsi/ginkgo v1.16.5 // indirect
-	github.com/onsi/gomega v1.13.0 // indirect
 	github.com/sirupsen/logrus v1.6.0
-	github.com/smartystreets/assertions v1.0.0 // indirect
 	github.com/spf13/cobra v1.0.0
 	github.com/spf13/viper v1.7.0
-	github.com/vulcanize/go-eth-state-node-iterator v1.0.0
-	go.uber.org/atomic v1.9.0 // indirect
-	go.uber.org/goleak v1.1.11 // indirect
-	go.uber.org/multierr v1.7.0 // indirect
-	go.uber.org/zap v1.19.1 // indirect
-	golang.org/x/crypto v0.0.0-20211209193657-4570a0811e8b // indirect
-	golang.org/x/lint v0.0.0-20200302205851-738671d3881b // indirect
-	golang.org/x/net v0.0.0-20211209124913-491a49abca63 // indirect
-	golang.org/x/sys v0.0.0-20211209171907-798191bca915 // indirect
-	golang.org/x/tools v0.1.8 // indirect
-	google.golang.org/appengine v1.6.6 // indirect
-	google.golang.org/protobuf v1.27.1 // indirect
-	lukechampine.com/blake3 v1.1.7 // indirect
+	github.com/vulcanize/go-eth-state-node-iterator v1.0.2
 )
 
-replace github.com/ethereum/go-ethereum v1.10.15 => github.com/vulcanize/go-ethereum v1.10.15-statediff-2.0.0
+replace github.com/ethereum/go-ethereum v1.10.17 => github.com/vulcanize/go-ethereum v1.10.17-statediff-4.0.1-alpha

main.go

@@ -18,7 +18,7 @@ package main
 import (
 	"github.com/sirupsen/logrus"
 
-	"github.com/vulcanize/eth-pg-ipfs-state-snapshot/cmd"
+	"github.com/vulcanize/ipld-eth-state-snapshot/cmd"
 )
 
 func main() {

pkg/snapshot/config.go

@@ -16,8 +16,13 @@
 package snapshot
 
 import (
+	"fmt"
+	"time"
+
+	"github.com/sirupsen/logrus"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
 	ethNode "github.com/ethereum/go-ethereum/statediff/indexer/node"
-	"github.com/ethereum/go-ethereum/statediff/indexer/postgres"
 	"github.com/spf13/viper"
 )
 
@@ -30,40 +35,57 @@ const (
 	LVL_DB_PATH       = "LVL_DB_PATH"
 )
 
+// SnapshotMode specifies the snapshot data output method
+type SnapshotMode string
+
+const (
+	PgSnapshot   SnapshotMode = "postgres"
+	FileSnapshot SnapshotMode = "file"
+
+	defaultOutputDir = "./snapshot_output"
+)
+
 // Config contains params for both databases the service uses
 type Config struct {
-	DB  *DBConfig
-	Eth *EthConfig
-}
-
-// DBConfig is config parameters for DB.
-type DBConfig struct {
-	Node       ethNode.Info
-	URI        string
-	ConnConfig postgres.ConnectionConfig
+	Eth  *EthConfig
+	DB   *DBConfig
+	File *FileConfig
 }
 
 // EthConfig is config parameters for the chain.
 type EthConfig struct {
 	LevelDBPath   string
 	AncientDBPath string
+	NodeInfo      ethNode.Info
 }
 
-func NewConfig() *Config {
-	ret := &Config{&DBConfig{}, &EthConfig{}}
-	ret.Init()
-	return ret
+// DBConfig is config parameters for DB.
+type DBConfig struct {
+	URI        string
+	ConnConfig postgres.Config
+}
+
+type FileConfig struct {
+	OutputDir string
+}
+
+func NewConfig(mode SnapshotMode) (*Config, error) {
+	ret := &Config{
+		&EthConfig{},
+		&DBConfig{},
+		&FileConfig{},
+	}
+	return ret, ret.Init(mode)
 }
 
 // Init Initialises config
-func (c *Config) Init() {
-	c.DB.dbInit()
+func (c *Config) Init(mode SnapshotMode) error {
 	viper.BindEnv("ethereum.nodeID", ETH_NODE_ID)
 	viper.BindEnv("ethereum.clientName", ETH_CLIENT_NAME)
 	viper.BindEnv("ethereum.genesisBlock", ETH_GENESIS_BLOCK)
 	viper.BindEnv("ethereum.networkID", ETH_NETWORK_ID)
 
-	c.DB.Node = ethNode.Info{
+	c.Eth.NodeInfo = ethNode.Info{
 		ID:           viper.GetString("ethereum.nodeID"),
 		ClientName:   viper.GetString("ethereum.clientName"),
 		GenesisBlock: viper.GetString("ethereum.genesisBlock"),
@@ -76,29 +98,50 @@ func (c *Config) Init() {
 
 	c.Eth.AncientDBPath = viper.GetString("leveldb.ancient")
 	c.Eth.LevelDBPath = viper.GetString("leveldb.path")
+
+	switch mode {
+	case FileSnapshot:
+		c.File.Init()
+	case PgSnapshot:
+		c.DB.Init()
+	default:
+		return fmt.Errorf("no output mode specified")
+	}
+	return nil
 }
 
-func (c *DBConfig) dbInit() {
-	viper.BindEnv("database.name", postgres.DATABASE_NAME)
-	viper.BindEnv("database.hostname", postgres.DATABASE_HOSTNAME)
-	viper.BindEnv("database.port", postgres.DATABASE_PORT)
-	viper.BindEnv("database.user", postgres.DATABASE_USER)
-	viper.BindEnv("database.password", postgres.DATABASE_PASSWORD)
-	viper.BindEnv("database.maxIdle", postgres.DATABASE_MAX_IDLE_CONNECTIONS)
-	viper.BindEnv("database.maxOpen", postgres.DATABASE_MAX_OPEN_CONNECTIONS)
-	viper.BindEnv("database.maxLifetime", postgres.DATABASE_MAX_CONN_LIFETIME)
+func (c *DBConfig) Init() {
+	viper.BindEnv("database.name", "DATABASE_NAME")
+	viper.BindEnv("database.hostname", "DATABASE_HOSTNAME")
+	viper.BindEnv("database.port", "DATABASE_PORT")
+	viper.BindEnv("database.user", "DATABASE_USER")
+	viper.BindEnv("database.password", "DATABASE_PASSWORD")
+	viper.BindEnv("database.maxIdle", "DATABASE_MAX_IDLE_CONNECTIONS")
+	viper.BindEnv("database.maxOpen", "DATABASE_MAX_OPEN_CONNECTIONS")
+	viper.BindEnv("database.maxLifetime", "DATABASE_MAX_CONN_LIFETIME")
 
-	dbParams := postgres.ConnectionParams{}
+	dbParams := postgres.Config{}
 	// DB params
-	dbParams.Name = viper.GetString("database.name")
+	dbParams.DatabaseName = viper.GetString("database.name")
 	dbParams.Hostname = viper.GetString("database.hostname")
 	dbParams.Port = viper.GetInt("database.port")
-	dbParams.User = viper.GetString("database.user")
+	dbParams.Username = viper.GetString("database.user")
 	dbParams.Password = viper.GetString("database.password")
-
-	c.URI = postgres.DbConnectionString(dbParams)
 	// Connection config
-	c.ConnConfig.MaxIdle = viper.GetInt("database.maxIdle")
-	c.ConnConfig.MaxOpen = viper.GetInt("database.maxOpen")
-	c.ConnConfig.MaxLifetime = viper.GetInt("database.maxLifetime")
+	dbParams.MaxIdle = viper.GetInt("database.maxIdle")
+	dbParams.MaxConns = viper.GetInt("database.maxOpen")
+	dbParams.MaxConnLifetime = time.Duration(viper.GetInt("database.maxLifetime")) * time.Second
+
+	c.ConnConfig = dbParams
+	c.URI = dbParams.DbConnectionString()
+}
+
+func (c *FileConfig) Init() error {
+	viper.BindEnv("file.outputDir", "FILE_OUTPUT_DIR")
+	c.OutputDir = viper.GetString("file.outputDir")
+	if c.OutputDir == "" {
+		logrus.Infof("no output directory set, using default: %s", defaultOutputDir)
+		c.OutputDir = defaultOutputDir
+	}
+	return nil
 }

pkg/snapshot/file/publisher.go

@@ -0,0 +1,295 @@
+// Copyright © 2020 Vulcanize, Inc
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+package publisher
+
+import (
+	"encoding/csv"
+	"fmt"
+	"math/big"
+	"os"
+	"path/filepath"
+	"sync/atomic"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ipfs/go-cid"
+	blockstore "github.com/ipfs/go-ipfs-blockstore"
+	dshelp "github.com/ipfs/go-ipfs-ds-help"
+	"github.com/multiformats/go-multihash"
+	"github.com/sirupsen/logrus"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/ipld"
+	nodeinfo "github.com/ethereum/go-ethereum/statediff/indexer/node"
+	"github.com/ethereum/go-ethereum/statediff/indexer/shared"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
+)
+
+var _ snapt.Publisher = (*publisher)(nil)
+
+var (
+	// tables written once per block
+	perBlockTables = []*snapt.Table{
+		&snapt.TableIPLDBlock,
+		&snapt.TableNodeInfo,
+		&snapt.TableHeader,
+	}
+	// tables written during state iteration
+	perNodeTables = []*snapt.Table{
+		&snapt.TableIPLDBlock,
+		&snapt.TableStateNode,
+		&snapt.TableStorageNode,
+	}
+)
+
+const logInterval = 1 * time.Minute
+
+type publisher struct {
+	dir     string // dir containing output files
+	writers fileWriters
+
+	nodeInfo nodeinfo.Info
+
+	startTime          time.Time
+	currBatchSize      uint
+	stateNodeCounter   uint64
+	storageNodeCounter uint64
+	codeNodeCounter    uint64
+	txCounter          uint32
+}
+
+type fileWriter struct {
+	*csv.Writer
+}
+
+// fileWriters wraps the file writers for each output table
+type fileWriters map[string]fileWriter
+
+type fileTx struct{ fileWriters }
+
+func (tx fileWriters) Commit() error {
+	for _, w := range tx {
+		w.Flush()
+		if err := w.Error(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+func (fileWriters) Rollback() error { return nil } // TODO: delete the file?
+
+func newFileWriter(path string) (ret fileWriter, err error) {
+	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
+	if err != nil {
+		return
+	}
+	ret = fileWriter{csv.NewWriter(file)}
+	return
+}
+
+func (tx fileWriters) write(tbl *snapt.Table, args ...interface{}) error {
+	row := tbl.ToCsvRow(args...)
+	return tx[tbl.Name].Write(row)
+}
+
+func makeFileWriters(dir string, tables []*snapt.Table) (fileWriters, error) {
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return nil, err
+	}
+	writers := fileWriters{}
+	for _, tbl := range tables {
+		w, err := newFileWriter(TableFile(dir, tbl.Name))
+		if err != nil {
+			return nil, err
+		}
+		writers[tbl.Name] = w
+	}
+	return writers, nil
+}
+
+// NewPublisher creates a publisher which writes to per-table CSV files which can be imported
+// with the Postgres COPY command.
+// The output directory will be created if it does not exist.
+func NewPublisher(path string, node nodeinfo.Info) (*publisher, error) {
+	if err := os.MkdirAll(path, 0777); err != nil {
+		return nil, fmt.Errorf("unable to make MkdirAll for path: %s err: %s", path, err)
+	}
+	writers, err := makeFileWriters(path, perBlockTables)
+	if err != nil {
+		return nil, err
+	}
+	pub := &publisher{
+		writers:   writers,
+		dir:       path,
+		nodeInfo:  node,
+		startTime: time.Now(),
+	}
+	go pub.logNodeCounters()
+	return pub, nil
+}
+
+func TableFile(dir, name string) string { return filepath.Join(dir, name+".csv") }
+
+func (p *publisher) txDir(index uint32) string {
+	return filepath.Join(p.dir, fmt.Sprintf("%010d", index))
+}
+
+func (p *publisher) BeginTx() (snapt.Tx, error) {
+	index := atomic.AddUint32(&p.txCounter, 1) - 1
+	dir := p.txDir(index)
+	writers, err := makeFileWriters(dir, perNodeTables)
+	if err != nil {
+		return nil, err
+	}
+
+	return fileTx{writers}, nil
+}
+
+// PublishRaw derives a cid from raw bytes and provided codec and multihash type, and writes it to the db tx
+// returns the CID and blockstore prefixed multihash key
+func (tx fileWriters) publishRaw(codec uint64, raw []byte, height *big.Int) (cid, prefixedKey string, err error) {
+	c, err := ipld.RawdataToCid(codec, raw, multihash.KECCAK_256)
+	if err != nil {
+		return
+	}
+	cid = c.String()
+	prefixedKey, err = tx.publishIPLD(c, raw, height)
+	return
+}
+
+func (tx fileWriters) publishIPLD(c cid.Cid, raw []byte, height *big.Int) (string, error) {
+	dbKey := dshelp.MultihashToDsKey(c.Hash())
+	prefixedKey := blockstore.BlockPrefix.String() + dbKey.String()
+	return prefixedKey, tx.write(&snapt.TableIPLDBlock, height.String(), prefixedKey, raw)
+}
+
+// PublishHeader writes the header to the ipfs backing pg datastore and adds secondary
+// indexes in the header_cids table
+func (p *publisher) PublishHeader(header *types.Header) error {
+	headerNode, err := ipld.NewEthHeader(header)
+	if err != nil {
+		return err
+	}
+	if _, err = p.writers.publishIPLD(headerNode.Cid(), headerNode.RawData(), header.Number); err != nil {
+		return err
+	}
+
+	mhKey := shared.MultihashKeyFromCID(headerNode.Cid())
+	err = p.writers.write(&snapt.TableNodeInfo, p.nodeInfo.GenesisBlock, p.nodeInfo.NetworkID, p.nodeInfo.ID,
+		p.nodeInfo.ClientName, p.nodeInfo.ChainID)
+	if err != nil {
+		return err
+	}
+	err = p.writers.write(&snapt.TableHeader, header.Number.String(), header.Hash().Hex(), header.ParentHash.Hex(),
+		headerNode.Cid().String(), 0, p.nodeInfo.ID, 0, header.Root.Hex(), header.TxHash.Hex(),
+		header.ReceiptHash.Hex(), header.UncleHash.Hex(), header.Bloom.Bytes(), header.Time, mhKey,
+		0, header.Coinbase.String())
+	if err != nil {
+		return err
+	}
+	return p.writers.Commit()
+}
+
+// PublishStateNode writes the state node to the ipfs backing datastore and adds secondary indexes
+// in the state_cids table
+func (p *publisher) PublishStateNode(node *snapt.Node, headerID string, height *big.Int, snapTx snapt.Tx) error {
+	var stateKey string
+	if !snapt.IsNullHash(node.Key) {
+		stateKey = node.Key.Hex()
+	}
+
+	tx := snapTx.(fileTx)
+	stateCIDStr, mhKey, err := tx.publishRaw(ipld.MEthStateTrie, node.Value, height)
+	if err != nil {
+		return err
+	}
+
+	err = tx.write(&snapt.TableStateNode, height.String(), headerID, stateKey, stateCIDStr, node.Path,
+		node.NodeType, false, mhKey)
+	if err != nil {
+		return err
+	}
+	// increment state node counter.
+	atomic.AddUint64(&p.stateNodeCounter, 1)
+	// increment current batch size counter
+	p.currBatchSize += 2
+	return err
+}
+
+// PublishStorageNode writes the storage node to the ipfs backing pg datastore and adds secondary
+// indexes in the storage_cids table
+func (p *publisher) PublishStorageNode(node *snapt.Node, headerID string, height *big.Int, statePath []byte, snapTx snapt.Tx) error {
+	var storageKey string
+	if !snapt.IsNullHash(node.Key) {
+		storageKey = node.Key.Hex()
+	}
+
+	tx := snapTx.(fileTx)
+	storageCIDStr, mhKey, err := tx.publishRaw(ipld.MEthStorageTrie, node.Value, height)
+	if err != nil {
+		return err
+	}
+
+	err = tx.write(&snapt.TableStorageNode, height.String(), headerID, statePath, storageKey, storageCIDStr, node.Path,
+		node.NodeType, false, mhKey)
+	if err != nil {
+		return err
+	}
+	// increment storage node counter.
+	atomic.AddUint64(&p.storageNodeCounter, 1)
+	// increment current batch size counter
+	p.currBatchSize += 2
+	return nil
+}
+
+// PublishCode writes code to the ipfs backing pg datastore
+func (p *publisher) PublishCode(height *big.Int, codeHash common.Hash, codeBytes []byte, snapTx snapt.Tx) error {
+	// no codec for code, doesn't matter though since blockstore key is multihash-derived
+	mhKey, err := shared.MultihashKeyFromKeccak256(codeHash)
+	if err != nil {
+		return fmt.Errorf("error deriving multihash key from codehash: %v", err)
+	}
+
+	tx := snapTx.(fileTx)
+	if err = tx.write(&snapt.TableIPLDBlock, height.String(), mhKey, codeBytes); err != nil {
+		return fmt.Errorf("error publishing code IPLD: %v", err)
+	}
+	// increment code node counter.
+	atomic.AddUint64(&p.codeNodeCounter, 1)
+	p.currBatchSize++
+	return nil
+}
+
+func (p *publisher) PrepareTxForBatch(tx snapt.Tx, maxBatchSize uint) (snapt.Tx, error) {
+	return tx, nil
+}
+
+// logNodeCounters periodically logs the number of node processed.
+func (p *publisher) logNodeCounters() {
+	t := time.NewTicker(logInterval)
+	for range t.C {
+		p.printNodeCounters("progress")
+	}
+}
+
+func (p *publisher) printNodeCounters(msg string) {
+	logrus.WithFields(logrus.Fields{
+		"runtime":       time.Now().Sub(p.startTime).String(),
+		"state nodes":   atomic.LoadUint64(&p.stateNodeCounter),
+		"storage nodes": atomic.LoadUint64(&p.storageNodeCounter),
+		"code nodes":    atomic.LoadUint64(&p.codeNodeCounter),
+	}).Info(msg)
+}

pkg/snapshot/file/publisher_test.go

@@ -0,0 +1,133 @@
+package publisher
+
+import (
+	"context"
+	"encoding/csv"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/ipld"
+	"github.com/jackc/pgx/v4"
+
+	fixt "github.com/vulcanize/ipld-eth-state-snapshot/fixture"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
+	"github.com/vulcanize/ipld-eth-state-snapshot/test"
+)
+
+var (
+	pgConfig = test.DefaultPgConfig
+	nodeInfo = test.DefaultNodeInfo
+	// tables ordered according to fkey depedencies
+	allTables = []*snapt.Table{
+		&snapt.TableIPLDBlock,
+		&snapt.TableNodeInfo,
+		&snapt.TableHeader,
+		&snapt.TableStateNode,
+		&snapt.TableStorageNode,
+	}
+)
+
+func writeFiles(t *testing.T, dir string) *publisher {
+	pub, err := NewPublisher(dir, nodeInfo)
+	test.NoError(t, err)
+	test.NoError(t, pub.PublishHeader(&fixt.Block1_Header))
+	tx, err := pub.BeginTx()
+	test.NoError(t, err)
+
+	headerID := fixt.Block1_Header.Hash().String()
+	test.NoError(t, pub.PublishStateNode(&fixt.Block1_StateNode0, headerID, fixt.Block1_Header.Number, tx))
+
+	test.NoError(t, tx.Commit())
+	return pub
+}
+
+// verify that we can parse the csvs
+// TODO check actual data
+func verifyFileData(t *testing.T, path string, tbl *snapt.Table) {
+	file, err := os.Open(path)
+	test.NoError(t, err)
+	r := csv.NewReader(file)
+	test.NoError(t, err)
+	r.FieldsPerRecord = len(tbl.Columns)
+
+	for {
+		_, err := r.Read()
+		if err == io.EOF {
+			break
+		}
+		test.NoError(t, err)
+	}
+}
+
+func TestWriting(t *testing.T) {
+	dir := t.TempDir()
+	// tempdir like /tmp/TempFoo/001/, TempFoo defaults to 0700
+	test.NoError(t, os.Chmod(filepath.Dir(dir), 0755))
+
+	pub := writeFiles(t, dir)
+
+	for _, tbl := range perBlockTables {
+		verifyFileData(t, TableFile(pub.dir, tbl.Name), tbl)
+	}
+	for i := uint32(0); i < pub.txCounter; i++ {
+		for _, tbl := range perNodeTables {
+			verifyFileData(t, TableFile(pub.txDir(i), tbl.Name), tbl)
+		}
+	}
+}
+
+// Note: DB user requires role membership "pg_read_server_files"
+func TestPgCopy(t *testing.T) {
+	test.NeedsDB(t)
+
+	dir := t.TempDir()
+	test.NoError(t, os.Chmod(filepath.Dir(dir), 0755))
+	pub := writeFiles(t, dir)
+
+	ctx := context.Background()
+	conn, err := pgx.Connect(ctx, pgConfig.DbConnectionString())
+	test.NoError(t, err)
+
+	// clear existing test data
+	pgDeleteTable := `DELETE FROM %s`
+	for _, tbl := range allTables {
+		_, err = conn.Exec(ctx, fmt.Sprintf(pgDeleteTable, tbl.Name))
+		test.NoError(t, err)
+	}
+
+	// copy from files
+	pgCopyStatement := `COPY %s FROM '%s' CSV`
+	for _, tbl := range perBlockTables {
+		stm := fmt.Sprintf(pgCopyStatement, tbl.Name, TableFile(pub.dir, tbl.Name))
+		_, err = conn.Exec(ctx, stm)
+		test.NoError(t, err)
+	}
+	for i := uint32(0); i < pub.txCounter; i++ {
+		for _, tbl := range perNodeTables {
+			stm := fmt.Sprintf(pgCopyStatement, tbl.Name, TableFile(pub.txDir(i), tbl.Name))
+			_, err = conn.Exec(ctx, stm)
+			test.NoError(t, err)
+		}
+	}
+
+	// check header was successfully committed
+	pgQueryHeader := `SELECT cid, block_hash
+					  FROM eth.header_cids
+				      WHERE block_number = $1`
+	type res struct {
+		CID       string
+		BlockHash string
+	}
+	var header res
+	err = conn.QueryRow(ctx, pgQueryHeader, fixt.Block1_Header.Number.Uint64()).Scan(
+		&header.CID, &header.BlockHash)
+	test.NoError(t, err)
+
+	headerNode, err := ipld.NewEthHeader(&fixt.Block1_Header)
+	test.NoError(t, err)
+	test.ExpectEqual(t, headerNode.Cid().String(), header.CID)
+	test.ExpectEqual(t, fixt.Block1_Header.Hash().String(), header.BlockHash)
+}

pkg/snapshot/mock/publisher.go

@@ -1,64 +0,0 @@
-package mock
-
-import (
-	"fmt"
-	"strings"
-	"testing"
-
-	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core/types"
-	"github.com/jmoiron/sqlx"
-
-	"github.com/golang/mock/gomock"
-
-	mocks "github.com/vulcanize/eth-pg-ipfs-state-snapshot/mocks/snapshot"
-	snapt "github.com/vulcanize/eth-pg-ipfs-state-snapshot/pkg/types"
-)
-
-type MockPublisher struct {
-	*mocks.MockPublisher
-}
-
-func NewMockPublisher(t *testing.T) *MockPublisher {
-	ctl := gomock.NewController(t)
-	return &MockPublisher{mocks.NewMockPublisher(ctl)}
-}
-
-func dump(funcname string, xs ...interface{}) {
-	if true {
-		return
-	}
-	fmt.Printf(">> %s", funcname)
-	fmt.Printf(strings.Repeat(" %+v", len(xs))+"\n", xs...)
-}
-
-func (p *MockPublisher) PublishHeader(header *types.Header) (int64, error) {
-	// fmt.Printf("PublishHeader %+v\n", header)
-	dump("PublishHeader", header)
-	return p.MockPublisher.PublishHeader(header)
-}
-func (p *MockPublisher) PublishStateNode(node *snapt.Node, headerID int64, tx *sqlx.Tx) (int64, error) {
-	dump("PublishStateNode", node, headerID)
-	return p.MockPublisher.PublishStateNode(node, headerID, tx)
-
-}
-func (p *MockPublisher) PublishStorageNode(node *snapt.Node, stateID int64, tx *sqlx.Tx) error {
-	dump("PublishStorageNode", node, stateID)
-	return p.MockPublisher.PublishStorageNode(node, stateID, tx)
-}
-func (p *MockPublisher) PublishCode(codeHash common.Hash, codeBytes []byte, tx *sqlx.Tx) error {
-	dump("PublishCode", codeHash, codeBytes)
-	return p.MockPublisher.PublishCode(codeHash, codeBytes, tx)
-}
-func (p *MockPublisher) BeginTx() (*sqlx.Tx, error) {
-	dump("BeginTx")
-	return p.MockPublisher.BeginTx()
-}
-func (p *MockPublisher) CommitTx(tx *sqlx.Tx) error {
-	dump("CommitTx", tx)
-	return p.MockPublisher.CommitTx(tx)
-}
-func (p *MockPublisher) PrepareTxForBatch(tx *sqlx.Tx, batchSize uint) (*sqlx.Tx, error) {
-	dump("PrepareTxForBatch", tx, batchSize)
-	return p.MockPublisher.PrepareTxForBatch(tx, batchSize)
-}

pkg/snapshot/mock/util.go

@@ -1,7 +1,10 @@
 package mock
 
-import "fmt"
-import "github.com/golang/mock/gomock"
+import (
+	"fmt"
+
+	"github.com/golang/mock/gomock"
+)
 
 type anyOfMatcher struct {
 	values []interface{}

pkg/snapshot/pg/publisher.go

@@ -0,0 +1,241 @@
+// Copyright © 2020 Vulcanize, Inc
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+package pg
+
+import (
+	"context"
+	"fmt"
+	"math/big"
+	"sync/atomic"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ipfs/go-cid"
+	blockstore "github.com/ipfs/go-ipfs-blockstore"
+	dshelp "github.com/ipfs/go-ipfs-ds-help"
+	"github.com/multiformats/go-multihash"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql"
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
+	"github.com/ethereum/go-ethereum/statediff/indexer/ipld"
+	"github.com/ethereum/go-ethereum/statediff/indexer/shared"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
+)
+
+var _ snapt.Publisher = (*publisher)(nil)
+
+const logInterval = 1 * time.Minute
+
+// Publisher is wrapper around DB.
+type publisher struct {
+	db                 *postgres.DB
+	currBatchSize      uint
+	stateNodeCounter   uint64
+	storageNodeCounter uint64
+	codeNodeCounter    uint64
+	startTime          time.Time
+}
+
+// NewPublisher creates Publisher
+func NewPublisher(db *postgres.DB) *publisher {
+	return &publisher{
+		db:        db,
+		startTime: time.Now(),
+	}
+}
+
+type pubTx struct {
+	sql.Tx
+	callback func()
+}
+
+func (tx pubTx) Rollback() error { return tx.Tx.Rollback(context.Background()) }
+func (tx pubTx) Commit() error {
+	if tx.callback != nil {
+		defer tx.callback()
+	}
+	return tx.Tx.Commit(context.Background())
+}
+func (tx pubTx) Exec(sql string, args ...interface{}) (sql.Result, error) {
+	return tx.Tx.Exec(context.Background(), sql, args...)
+}
+
+func (p *publisher) BeginTx() (snapt.Tx, error) {
+	tx, err := p.db.Begin(context.Background())
+	if err != nil {
+		return nil, err
+	}
+	go p.logNodeCounters()
+	return pubTx{tx, func() {
+		p.printNodeCounters("final stats")
+	}}, nil
+}
+
+// PublishRaw derives a cid from raw bytes and provided codec and multihash type, and writes it to the db tx
+// returns the CID and blockstore prefixed multihash key
+func (tx pubTx) publishRaw(codec uint64, raw []byte, height *big.Int) (cid, prefixedKey string, err error) {
+	c, err := ipld.RawdataToCid(codec, raw, multihash.KECCAK_256)
+	if err != nil {
+		return
+	}
+	cid = c.String()
+	prefixedKey, err = tx.publishIPLD(c, raw, height)
+	return
+}
+
+func (tx pubTx) publishIPLD(c cid.Cid, raw []byte, height *big.Int) (string, error) {
+	dbKey := dshelp.MultihashToDsKey(c.Hash())
+	prefixedKey := blockstore.BlockPrefix.String() + dbKey.String()
+	_, err := tx.Exec(snapt.TableIPLDBlock.ToInsertStatement(), height.Uint64(), prefixedKey, raw)
+	return prefixedKey, err
+}
+
+// PublishHeader writes the header to the ipfs backing pg datastore and adds secondary indexes in the header_cids table
+func (p *publisher) PublishHeader(header *types.Header) (err error) {
+	headerNode, err := ipld.NewEthHeader(header)
+	if err != nil {
+		return err
+	}
+
+	snapTx, err := p.db.Begin(context.Background())
+	if err != nil {
+		return err
+	}
+	tx := pubTx{snapTx, nil}
+	defer func() { err = snapt.CommitOrRollback(tx, err) }()
+
+	if _, err = tx.publishIPLD(headerNode.Cid(), headerNode.RawData(), header.Number); err != nil {
+		return err
+	}
+
+	mhKey := shared.MultihashKeyFromCID(headerNode.Cid())
+	_, err = tx.Exec(snapt.TableHeader.ToInsertStatement(), header.Number.Uint64(), header.Hash().Hex(),
+		header.ParentHash.Hex(), headerNode.Cid().String(), "0", p.db.NodeID(), "0",
+		header.Root.Hex(), header.TxHash.Hex(), header.ReceiptHash.Hex(), header.UncleHash.Hex(),
+		header.Bloom.Bytes(), header.Time, mhKey, 0, header.Coinbase.String())
+	return err
+}
+
+// PublishStateNode writes the state node to the ipfs backing datastore and adds secondary indexes in the state_cids table
+func (p *publisher) PublishStateNode(node *snapt.Node, headerID string, height *big.Int, snapTx snapt.Tx) error {
+	var stateKey string
+	if !snapt.IsNullHash(node.Key) {
+		stateKey = node.Key.Hex()
+	}
+
+	tx := snapTx.(pubTx)
+	stateCIDStr, mhKey, err := tx.publishRaw(ipld.MEthStateTrie, node.Value, height)
+	if err != nil {
+		return err
+	}
+
+	_, err = tx.Exec(snapt.TableStateNode.ToInsertStatement(),
+		height.Uint64(), headerID, stateKey, stateCIDStr, node.Path, node.NodeType, false, mhKey)
+	if err != nil {
+		return err
+	}
+	// increment state node counter.
+	atomic.AddUint64(&p.stateNodeCounter, 1)
+
+	// increment current batch size counter
+	p.currBatchSize += 2
+	return err
+}
+
+// PublishStorageNode writes the storage node to the ipfs backing pg datastore and adds secondary indexes in the storage_cids table
+func (p *publisher) PublishStorageNode(node *snapt.Node, headerID string, height *big.Int, statePath []byte, snapTx snapt.Tx) error {
+	var storageKey string
+	if !snapt.IsNullHash(node.Key) {
+		storageKey = node.Key.Hex()
+	}
+
+	tx := snapTx.(pubTx)
+	storageCIDStr, mhKey, err := tx.publishRaw(ipld.MEthStorageTrie, node.Value, height)
+	if err != nil {
+		return err
+	}
+
+	_, err = tx.Exec(snapt.TableStorageNode.ToInsertStatement(),
+		height.Uint64(), headerID, statePath, storageKey, storageCIDStr, node.Path, node.NodeType, false, mhKey)
+	if err != nil {
+		return err
+	}
+	// increment storage node counter.
+	atomic.AddUint64(&p.storageNodeCounter, 1)
+
+	// increment current batch size counter
+	p.currBatchSize += 2
+	return err
+}
+
+// PublishCode writes code to the ipfs backing pg datastore
+func (p *publisher) PublishCode(height *big.Int, codeHash common.Hash, codeBytes []byte, snapTx snapt.Tx) error {
+	// no codec for code, doesn't matter though since blockstore key is multihash-derived
+	mhKey, err := shared.MultihashKeyFromKeccak256(codeHash)
+	if err != nil {
+		return fmt.Errorf("error deriving multihash key from codehash: %v", err)
+	}
+
+	tx := snapTx.(pubTx)
+	if _, err = tx.Exec(snapt.TableIPLDBlock.ToInsertStatement(), height.Uint64(), mhKey, codeBytes); err != nil {
+		return fmt.Errorf("error publishing code IPLD: %v", err)
+	}
+
+	// increment code node counter.
+	atomic.AddUint64(&p.codeNodeCounter, 1)
+
+	p.currBatchSize++
+	return nil
+}
+
+func (p *publisher) PrepareTxForBatch(tx snapt.Tx, maxBatchSize uint) (snapt.Tx, error) {
+	var err error
+	// maximum batch size reached, commit the current transaction and begin a new transaction.
+	if maxBatchSize <= p.currBatchSize {
+		if err = tx.Commit(); err != nil {
+			return nil, err
+		}
+
+		snapTx, err := p.db.Begin(context.Background())
+		tx = pubTx{Tx: snapTx}
+		if err != nil {
+			return nil, err
+		}
+
+		p.currBatchSize = 0
+	}
+
+	return tx, nil
+}
+
+// logNodeCounters periodically logs the number of node processed.
+func (p *publisher) logNodeCounters() {
+	t := time.NewTicker(logInterval)
+	for range t.C {
+		p.printNodeCounters("progress")
+	}
+}
+
+func (p *publisher) printNodeCounters(msg string) {
+	log.WithFields(log.Fields{
+		"runtime":       time.Now().Sub(p.startTime).String(),
+		"state nodes":   atomic.LoadUint64(&p.stateNodeCounter),
+		"storage nodes": atomic.LoadUint64(&p.storageNodeCounter),
+		"code nodes":    atomic.LoadUint64(&p.codeNodeCounter),
+	}).Info(msg)
+}

pkg/snapshot/pg/publisher_test.go

@@ -0,0 +1,79 @@
+package pg
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
+	"github.com/ethereum/go-ethereum/statediff/indexer/ipld"
+	"github.com/jackc/pgx/v4"
+
+	fixt "github.com/vulcanize/ipld-eth-state-snapshot/fixture"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
+	"github.com/vulcanize/ipld-eth-state-snapshot/test"
+)
+
+var (
+	pgConfig = test.DefaultPgConfig
+	nodeInfo = test.DefaultNodeInfo
+	// tables ordered according to fkey depedencies
+	allTables = []*snapt.Table{
+		&snapt.TableIPLDBlock,
+		&snapt.TableNodeInfo,
+		&snapt.TableHeader,
+		&snapt.TableStateNode,
+		&snapt.TableStorageNode,
+	}
+)
+
+func writeData(t *testing.T) *publisher {
+	driver, err := postgres.NewPGXDriver(context.Background(), pgConfig, nodeInfo)
+	test.NoError(t, err)
+	pub := NewPublisher(postgres.NewPostgresDB(driver))
+	test.NoError(t, pub.PublishHeader(&fixt.Block1_Header))
+	tx, err := pub.BeginTx()
+	test.NoError(t, err)
+
+	headerID := fixt.Block1_Header.Hash().String()
+	test.NoError(t, pub.PublishStateNode(&fixt.Block1_StateNode0, headerID, fixt.Block1_Header.Number, tx))
+
+	test.NoError(t, tx.Commit())
+	return pub
+}
+
+// Note: DB user requires role membership "pg_read_server_files"
+func TestBasic(t *testing.T) {
+	test.NeedsDB(t)
+
+	ctx := context.Background()
+	conn, err := pgx.Connect(ctx, pgConfig.DbConnectionString())
+	test.NoError(t, err)
+
+	// clear existing test data
+	pgDeleteTable := `DELETE FROM %s`
+	for _, tbl := range allTables {
+		_, err = conn.Exec(ctx, fmt.Sprintf(pgDeleteTable, tbl.Name))
+		test.NoError(t, err)
+	}
+
+	_ = writeData(t)
+
+	// check header was successfully committed
+	pgQueryHeader := `SELECT cid, block_hash
+					  FROM eth.header_cids
+				      WHERE block_number = $1`
+	type res struct {
+		CID       string
+		BlockHash string
+	}
+	var header res
+	err = conn.QueryRow(ctx, pgQueryHeader, fixt.Block1_Header.Number.Uint64()).Scan(
+		&header.CID, &header.BlockHash)
+	test.NoError(t, err)
+
+	headerNode, err := ipld.NewEthHeader(&fixt.Block1_Header)
+	test.NoError(t, err)
+	test.ExpectEqual(t, headerNode.Cid().String(), header.CID)
+	test.ExpectEqual(t, fixt.Block1_Header.Hash().String(), header.BlockHash)
+}

pkg/snapshot/publisher.go

@@ -1,215 +0,0 @@
-// Copyright © 2020 Vulcanize, Inc
-//
-// This program is free software: you can redistribute it and/or modify
-// it under the terms of the GNU Affero General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-//
-// This program is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-// GNU Affero General Public License for more details.
-//
-// You should have received a copy of the GNU Affero General Public License
-// along with this program. If not, see <http://www.gnu.org/licenses/>.
-
-package snapshot
-
-import (
-	"bytes"
-	"fmt"
-	"sync/atomic"
-	"time"
-
-	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core/types"
-	"github.com/jmoiron/sqlx"
-	"github.com/multiformats/go-multihash"
-	"github.com/sirupsen/logrus"
-
-	"github.com/ethereum/go-ethereum/statediff/indexer/ipfs/ipld"
-	"github.com/ethereum/go-ethereum/statediff/indexer/postgres"
-	"github.com/ethereum/go-ethereum/statediff/indexer/shared"
-	snapt "github.com/vulcanize/eth-pg-ipfs-state-snapshot/pkg/types"
-)
-
-const logInterval = 1 * time.Minute
-
-// Publisher is wrapper around DB.
-type publisher struct {
-	db                 *postgres.DB
-	currBatchSize      uint
-	stateNodeCounter   uint64
-	storageNodeCounter uint64
-	codeNodeCounter    uint64
-	startTime          time.Time
-}
-
-// NewPublisher creates Publisher
-func NewPublisher(db *postgres.DB) *publisher {
-	return &publisher{
-		db:            db,
-		currBatchSize: 0,
-		startTime:     time.Now(),
-	}
-}
-
-func (p *publisher) BeginTx() (*sqlx.Tx, error) {
-	tx, err := p.db.Beginx()
-	if err != nil {
-		return nil, err
-	}
-	go p.logNodeCounters()
-	return tx, nil
-}
-
-func (p *publisher) CommitTx(tx *sqlx.Tx) error {
-	logrus.Info("----- final counts -----")
-	p.printNodeCounters()
-	return tx.Commit()
-}
-
-// PublishHeader writes the header to the ipfs backing pg datastore and adds secondary indexes in the header_cids table
-func (p *publisher) PublishHeader(header *types.Header) (int64, error) {
-	headerNode, err := ipld.NewEthHeader(header)
-	if err != nil {
-		return 0, err
-	}
-
-	tx, err := p.db.Beginx()
-	if err != nil {
-		return 0, err
-	}
-
-	defer func() {
-		if p := recover(); p != nil {
-			shared.Rollback(tx)
-			panic(p)
-		} else if err != nil {
-			shared.Rollback(tx)
-		} else {
-			err = tx.Commit()
-		}
-	}()
-
-	if err = shared.PublishIPLD(tx, headerNode); err != nil {
-		return 0, err
-	}
-
-	mhKey, _ := shared.MultihashKeyFromCIDString(headerNode.Cid().String())
-	var headerID int64
-	err = tx.QueryRowx(`INSERT INTO eth.header_cids (block_number, block_hash, parent_hash, cid, td, node_id, reward, state_root, tx_root, receipt_root, uncle_root, bloom, timestamp, mh_key, times_validated)
- 								VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
- 								ON CONFLICT (block_number, block_hash) DO UPDATE SET block_number = header_cids.block_number
- 								RETURNING id`,
-		header.Number.Uint64(), header.Hash().Hex(), header.ParentHash.Hex(), headerNode.Cid().String(), "0", p.db.NodeID, "0", header.Root.Hex(), header.TxHash.Hex(),
-		header.ReceiptHash.Hex(), header.UncleHash.Hex(), header.Bloom.Bytes(), header.Time, mhKey, 0).Scan(&headerID)
-
-	return headerID, err
-}
-
-// PublishStateNode writes the state node to the ipfs backing datastore and adds secondary indexes in the state_cids table
-func (p *publisher) PublishStateNode(node *snapt.Node, headerID int64, tx *sqlx.Tx) (int64, error) {
-	var stateID int64
-	var stateKey string
-	if !bytes.Equal(node.Key.Bytes(), nullHash.Bytes()) {
-		stateKey = node.Key.Hex()
-	}
-
-	stateCIDStr, mhKey, err := shared.PublishRaw(tx, ipld.MEthStateTrie, multihash.KECCAK_256, node.Value)
-	if err != nil {
-		return 0, err
-	}
-
-	err = tx.QueryRowx(`INSERT INTO eth.state_cids (header_id, state_leaf_key, cid, state_path, node_type, diff, mh_key) VALUES ($1, $2, $3, $4, $5, $6, $7)
- 									ON CONFLICT (header_id, state_path) DO UPDATE SET (state_leaf_key, cid, node_type, diff, mh_key) = ($2, $3, $5, $6, $7)
- 									RETURNING id`,
-		headerID, stateKey, stateCIDStr, node.Path, node.NodeType, false, mhKey).Scan(&stateID)
-
-	// increment state node counter.
-	atomic.AddUint64(&p.stateNodeCounter, 1)
-
-	// increment current batch size counter
-	p.currBatchSize += 2
-	return stateID, err
-}
-
-// PublishStorageNode writes the storage node to the ipfs backing pg datastore and adds secondary indexes in the storage_cids table
-func (p *publisher) PublishStorageNode(node *snapt.Node, stateID int64, tx *sqlx.Tx) error {
-	var storageKey string
-	if !bytes.Equal(node.Key.Bytes(), nullHash.Bytes()) {
-		storageKey = node.Key.Hex()
-	}
-
-	storageCIDStr, mhKey, err := shared.PublishRaw(tx, ipld.MEthStorageTrie, multihash.KECCAK_256, node.Value)
-	if err != nil {
-		return err
-	}
-
-	_, err = tx.Exec(`INSERT INTO eth.storage_cids (state_id, storage_leaf_key, cid, storage_path, node_type, diff, mh_key) VALUES ($1, $2, $3, $4, $5, $6, $7)
-                              	ON CONFLICT (state_id, storage_path) DO UPDATE SET (storage_leaf_key, cid, node_type, diff, mh_key) = ($2, $3, $5, $6, $7)`,
-		stateID, storageKey, storageCIDStr, node.Path, node.NodeType, false, mhKey)
-	if err != nil {
-		return err
-	}
-
-	// increment storage node counter.
-	atomic.AddUint64(&p.storageNodeCounter, 1)
-
-	// increment current batch size counter
-	p.currBatchSize += 2
-	return nil
-}
-
-// PublishCode writes code to the ipfs backing pg datastore
-func (p *publisher) PublishCode(codeHash common.Hash, codeBytes []byte, tx *sqlx.Tx) error {
-	// no codec for code, doesn't matter though since blockstore key is multihash-derived
-	mhKey, err := shared.MultihashKeyFromKeccak256(codeHash)
-	if err != nil {
-		return fmt.Errorf("error deriving multihash key from codehash: %v", err)
-	}
-
-	if err = shared.PublishDirect(tx, mhKey, codeBytes); err != nil {
-		return fmt.Errorf("error publishing code IPLD: %v", err)
-	}
-
-	// increment code node counter.
-	atomic.AddUint64(&p.codeNodeCounter, 1)
-
-	p.currBatchSize++
-	return nil
-}
-
-func (p *publisher) PrepareTxForBatch(tx *sqlx.Tx, maxBatchSize uint) (*sqlx.Tx, error) {
-	var err error
-	// maximum batch size reached, commit the current transaction and begin a new transaction.
-	if maxBatchSize <= p.currBatchSize {
-		if err = tx.Commit(); err != nil {
-			return nil, err
-		}
-
-		tx, err = p.db.Beginx()
-		if err != nil {
-			return nil, err
-		}
-
-		p.currBatchSize = 0
-	}
-
-	return tx, nil
-}
-
-// logNodeCounters periodically logs the number of node processed.
-func (p *publisher) logNodeCounters() {
-	t := time.NewTicker(logInterval)
-	for range t.C {
-		p.printNodeCounters()
-	}
-}
-
-func (p *publisher) printNodeCounters() {
-	logrus.Infof("runtime: %s", time.Now().Sub(p.startTime).String())
-	logrus.Infof("processed state nodes: %d", atomic.LoadUint64(&p.stateNodeCounter))
-	logrus.Infof("processed storage nodes: %d", atomic.LoadUint64(&p.storageNodeCounter))
-	logrus.Infof("processed code nodes: %d", atomic.LoadUint64(&p.codeNodeCounter))
-}

pkg/snapshot/service.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"math/big"
 	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -28,19 +29,15 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/rlp"
-	"github.com/ethereum/go-ethereum/statediff/indexer/postgres"
-	"github.com/ethereum/go-ethereum/statediff/indexer/shared"
 	"github.com/ethereum/go-ethereum/trie"
-	"github.com/jmoiron/sqlx"
-	"github.com/sirupsen/logrus"
+	log "github.com/sirupsen/logrus"
 
-	. "github.com/vulcanize/eth-pg-ipfs-state-snapshot/pkg/types"
 	iter "github.com/vulcanize/go-eth-state-node-iterator"
+	. "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
 )
 
 var (
-	nullHash          = common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000000")
-	emptyNode, _      = rlp.EncodeToBytes([]byte{})
+	emptyNode, _      = rlp.EncodeToBytes(&[]byte{})
 	emptyCodeHash     = crypto.Keccak256([]byte{})
 	emptyContractRoot = crypto.Keccak256Hash(emptyNode)
 
@@ -54,29 +51,30 @@ type Service struct {
 	stateDB       state.Database
 	ipfsPublisher Publisher
 	maxBatchSize  uint
-}
-
-func NewPostgresDB(con *DBConfig) (*postgres.DB, error) {
-	pgDB, err := postgres.NewDB(con.URI, con.ConnConfig, con.Node)
-	if err != nil {
-		return nil, err
-	}
-	return pgDB, nil
+	tracker       iteratorTracker
+	recoveryFile  string
 }
 
 func NewLevelDB(con *EthConfig) (ethdb.Database, error) {
-	return rawdb.NewLevelDBDatabaseWithFreezer(
-		con.LevelDBPath, 1024, 256, con.AncientDBPath, "eth-pg-ipfs-state-snapshot", false,
+	println(con.LevelDBPath)
+	println(con.AncientDBPath)
+	edb, err := rawdb.NewLevelDBDatabaseWithFreezer(
+		con.LevelDBPath, 1024, 256, con.AncientDBPath, "ipld-eth-state-snapshot", true,
 	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create NewLevelDBDatabaseWithFreezer: %s", err)
+	}
+	return edb, nil
 }
 
 // NewSnapshotService creates Service.
-func NewSnapshotService(edb ethdb.Database, pub Publisher) (*Service, error) {
+func NewSnapshotService(edb ethdb.Database, pub Publisher, recoveryFile string) (*Service, error) {
 	return &Service{
 		ethDB:         edb,
 		stateDB:       state.NewDatabase(edb),
 		ipfsPublisher: pub,
 		maxBatchSize:  defaultBatchSize,
+		recoveryFile:  recoveryFile,
 	}, nil
 }
 
@@ -88,35 +86,70 @@ type SnapshotParams struct {
 func (s *Service) CreateSnapshot(params SnapshotParams) error {
 	// extract header from lvldb and publish to PG-IPFS
 	// hold onto the headerID so that we can link the state nodes to this header
-	logrus.Infof("Creating snapshot at height %d", params.Height)
+	log.Infof("Creating snapshot at height %d", params.Height)
 	hash := rawdb.ReadCanonicalHash(s.ethDB, params.Height)
 	header := rawdb.ReadHeader(s.ethDB, hash, params.Height)
 	if header == nil {
 		return fmt.Errorf("unable to read canonical header at height %d", params.Height)
 	}
 
-	logrus.Infof("head hash: %s head height: %d", hash.Hex(), params.Height)
+	log.Infof("head hash: %s head height: %d", hash.Hex(), params.Height)
 
-	headerID, err := s.ipfsPublisher.PublishHeader(header)
+	err := s.ipfsPublisher.PublishHeader(header)
 	if err != nil {
 		return err
 	}
 
-	t, err := s.stateDB.OpenTrie(header.Root)
+	tree, err := s.stateDB.OpenTrie(header.Root)
 	if err != nil {
 		return err
 	}
-	if params.Workers > 0 {
-		return s.createSnapshotAsync(t, headerID, params.Workers)
+	headerID := header.Hash().String()
+	s.tracker = newTracker(s.recoveryFile, int(params.Workers))
+	go s.tracker.run()
+	go s.tracker.captureSignal()
+
+	var iters []trie.NodeIterator
+	// attempt to restore from recovery file if it exists
+	iters, err = s.tracker.restore(tree)
+	if err != nil {
+		return err
+	}
+	if iters != nil {
+		if params.Workers < uint(len(iters)) {
+			return fmt.Errorf(
+				"number of recovered workers (%d) is greater than number configured (%d)",
+				len(iters), params.Workers,
+			)
+		}
+	} else { // nothing to restore
+		if params.Workers > 1 {
+			iters = iter.SubtrieIterators(tree, params.Workers)
+		} else {
+			iters = []trie.NodeIterator{tree.NodeIterator(nil)}
+		}
+		for i, it := range iters {
+			iters[i] = s.tracker.tracked(it)
+		}
+	}
+
+	defer func() {
+		err := s.tracker.haltAndDump()
+		if err != nil {
+			log.Error("failed to write recovery file: ", err)
+		}
+	}()
+
+	if len(iters) > 0 {
+		return s.createSnapshotAsync(iters, headerID, new(big.Int).SetUint64(params.Height))
 	} else {
-		return s.createSnapshot(t.NodeIterator(nil), headerID)
+		return s.createSnapshot(iters[0], headerID, new(big.Int).SetUint64(params.Height))
 	}
-	return nil
 }
 
 // Create snapshot up to head (ignores height param)
 func (s *Service) CreateLatestSnapshot(workers uint) error {
-	logrus.Info("Creating snapshot at head")
+	log.Info("Creating snapshot at head")
 	hash := rawdb.ReadHeadHeaderHash(s.ethDB)
 	height := rawdb.ReadHeaderNumber(s.ethDB, hash)
 	if height == nil {
@@ -135,7 +168,7 @@ func resolveNode(it trie.NodeIterator, trieDB *trie.Database) (*nodeResult, erro
 	if it.Leaf() {
 		return nil, nil
 	}
-	if bytes.Equal(nullHash.Bytes(), it.Hash().Bytes()) {
+	if IsNullHash(it.Hash()) {
 		return nil, nil
 	}
 
@@ -163,22 +196,12 @@ func resolveNode(it trie.NodeIterator, trieDB *trie.Database) (*nodeResult, erro
 	}, nil
 }
 
-func (s *Service) createSnapshot(it trie.NodeIterator, headerID int64) error {
+func (s *Service) createSnapshot(it trie.NodeIterator, headerID string, height *big.Int) error {
 	tx, err := s.ipfsPublisher.BeginTx()
 	if err != nil {
 		return err
 	}
-
-	defer func() {
-		if rec := recover(); rec != nil {
-			shared.Rollback(tx)
-			panic(rec)
-		} else if err != nil {
-			shared.Rollback(tx)
-		} else {
-			err = s.ipfsPublisher.CommitTx(tx)
-		}
-	}()
+	defer func() { err = CommitOrRollback(tx, err) }()
 
 	for it.Next(true) {
 		res, err := resolveNode(it, s.stateDB.TrieDB())
@@ -196,7 +219,8 @@ func (s *Service) createSnapshot(it trie.NodeIterator, headerID int64) error {
 
 		switch res.node.NodeType {
 		case Leaf:
-			// if the node is a leaf, decode the account and publish the associated storage trie nodes if there are any
+			// if the node is a leaf, decode the account and publish the associated storage trie
+			// nodes if there are any
 			var account types.StateAccount
 			if err := rlp.DecodeBytes(res.elements[1].([]byte), &account); err != nil {
 				return fmt.Errorf(
@@ -207,7 +231,7 @@ func (s *Service) createSnapshot(it trie.NodeIterator, headerID int64) error {
 			encodedPath := trie.HexToCompact(valueNodePath)
 			leafKey := encodedPath[1:]
 			res.node.Key = common.BytesToHash(leafKey)
-			stateID, err := s.ipfsPublisher.PublishStateNode(&res.node, headerID, tx)
+			err := s.ipfsPublisher.PublishStateNode(&res.node, headerID, height, tx)
 			if err != nil {
 				return err
 			}
@@ -217,57 +241,60 @@ func (s *Service) createSnapshot(it trie.NodeIterator, headerID int64) error {
 				codeHash := common.BytesToHash(account.CodeHash)
 				codeBytes := rawdb.ReadCode(s.ethDB, codeHash)
 				if len(codeBytes) == 0 {
-					logrus.Error("Code is missing", "account", common.BytesToHash(it.LeafKey()))
+					log.Error("Code is missing", "account", common.BytesToHash(it.LeafKey()))
 					return errors.New("missing code")
 				}
 
-				if err = s.ipfsPublisher.PublishCode(codeHash, codeBytes, tx); err != nil {
+				if err = s.ipfsPublisher.PublishCode(height, codeHash, codeBytes, tx); err != nil {
 					return err
 				}
 			}
 
-			if tx, err = s.storageSnapshot(account.Root, stateID, tx); err != nil {
+			if tx, err = s.storageSnapshot(account.Root, headerID, height, res.node.Path, tx); err != nil {
 				return fmt.Errorf("failed building storage snapshot for account %+v\r\nerror: %w", account, err)
 			}
 		case Extension, Branch:
 			res.node.Key = common.BytesToHash([]byte{})
-			if _, err := s.ipfsPublisher.PublishStateNode(&res.node, headerID, tx); err != nil {
+			if err := s.ipfsPublisher.PublishStateNode(&res.node, headerID, height, tx); err != nil {
 				return err
 			}
 		default:
 			return errors.New("unexpected node type")
 		}
-		return nil
 	}
 	return it.Error()
 }
 
 // Full-trie concurrent snapshot
-func (s *Service) createSnapshotAsync(tree state.Trie, headerID int64, workers uint) error {
+func (s *Service) createSnapshotAsync(iters []trie.NodeIterator, headerID string, height *big.Int) error {
 	errors := make(chan error)
 	var wg sync.WaitGroup
-	for _, it := range iter.SubtrieIterators(tree, workers) {
+	for _, it := range iters {
 		wg.Add(1)
-		go func() {
+		go func(it trie.NodeIterator) {
 			defer wg.Done()
-			if err := s.createSnapshot(it, headerID); err != nil {
+			if err := s.createSnapshot(it, headerID, height); err != nil {
 				errors <- err
 			}
-		}()
+		}(it)
 	}
+
+	done := make(chan struct{})
 	go func() {
-		defer close(errors)
 		wg.Wait()
+		done <- struct{}{}
 	}()
 
+	var err error
 	select {
-	case err := <-errors:
-		return err
+	case err = <-errors:
+	case <-done:
+		close(errors)
 	}
-	return nil
+	return err
 }
 
-func (s *Service) storageSnapshot(sr common.Hash, stateID int64, tx *sqlx.Tx) (*sqlx.Tx, error) {
+func (s *Service) storageSnapshot(sr common.Hash, headerID string, height *big.Int, statePath []byte, tx Tx) (Tx, error) {
 	if bytes.Equal(sr.Bytes(), emptyContractRoot.Bytes()) {
 		return tx, nil
 	}
@@ -311,7 +338,7 @@ func (s *Service) storageSnapshot(sr common.Hash, stateID int64, tx *sqlx.Tx) (*
 		default:
 			return nil, errors.New("unexpected node type")
 		}
-		if err = s.ipfsPublisher.PublishStorageNode(&res.node, stateID, tx); err != nil {
+		if err = s.ipfsPublisher.PublishStorageNode(&res.node, headerID, height, statePath, tx); err != nil {
 			return nil, err
 		}
 	}

pkg/snapshot/service_test.go

@@ -1,99 +1,141 @@
 package snapshot
 
 import (
-	"fmt"
+	"errors"
+	"math/big"
 	"os"
 	"path/filepath"
 	"testing"
 
 	"github.com/golang/mock/gomock"
 
-	ethNode "github.com/ethereum/go-ethereum/statediff/indexer/node"
-	"github.com/ethereum/go-ethereum/statediff/indexer/postgres"
-
-	fixt "github.com/vulcanize/eth-pg-ipfs-state-snapshot/fixture"
-	"github.com/vulcanize/eth-pg-ipfs-state-snapshot/pkg/snapshot/mock"
+	fixt "github.com/vulcanize/ipld-eth-state-snapshot/fixture"
+	mock "github.com/vulcanize/ipld-eth-state-snapshot/mocks/snapshot"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
+	"github.com/vulcanize/ipld-eth-state-snapshot/test"
 )
 
 func testConfig(leveldbpath, ancientdbpath string) *Config {
-	dbParams := postgres.ConnectionParams{
-		Name:     "snapshot_test",
-		Hostname: "localhost",
-		Port:     5432,
-		User:     "tester",
-		Password: "test_pw",
-	}
-	connconfig := postgres.ConnectionConfig{
-		MaxIdle:     0,
-		MaxLifetime: 0,
-		MaxOpen:     4,
-	}
-	nodeinfo := ethNode.Info{
-		ID:           "eth_node_id",
-		ClientName:   "eth_client",
-		GenesisBlock: "X",
-		NetworkID:    "eth_network",
-		ChainID:      0,
-	}
-
 	return &Config{
-		DB: &DBConfig{
-			Node:       nodeinfo,
-			URI:        postgres.DbConnectionString(dbParams),
-			ConnConfig: connconfig,
-		},
 		Eth: &EthConfig{
 			LevelDBPath:   leveldbpath,
 			AncientDBPath: ancientdbpath,
+			NodeInfo:      test.DefaultNodeInfo,
+		},
+		DB: &DBConfig{
+			URI:        test.DefaultPgConfig.DbConnectionString(),
+			ConnConfig: test.DefaultPgConfig,
 		},
 	}
 }
 
+func makeMocks(t *testing.T) (*mock.MockPublisher, *mock.MockTx) {
+	ctl := gomock.NewController(t)
+	pub := mock.NewMockPublisher(ctl)
+	tx := mock.NewMockTx(ctl)
+	return pub, tx
+}
+
 func TestCreateSnapshot(t *testing.T) {
-	wd, err := os.Getwd()
-	if err != nil {
-		panic(err)
+	runCase := func(t *testing.T, workers int) {
+		pub, tx := makeMocks(t)
+		pub.EXPECT().PublishHeader(gomock.Eq(&fixt.Block1_Header))
+		pub.EXPECT().BeginTx().Return(tx, nil).
+			Times(workers)
+		pub.EXPECT().PrepareTxForBatch(gomock.Any(), gomock.Any()).Return(tx, nil).
+			AnyTimes()
+		pub.EXPECT().PublishStateNode(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+			Times(len(fixt.Block1_StateNodePaths))
+
+		// TODO: fixtures for storage node
+		// pub.EXPECT().PublishStorageNode(gomock.Eq(fixt.StorageNode), gomock.Eq(int64(0)), gomock.Any())
+
+		tx.EXPECT().Commit().
+			Times(workers)
+
+		config := testConfig(fixt.ChaindataPath, fixt.AncientdataPath)
+		edb, err := NewLevelDB(config.Eth)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer edb.Close()
+
+		recovery := filepath.Join(t.TempDir(), "recover.csv")
+		service, err := NewSnapshotService(edb, pub, recovery)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		params := SnapshotParams{Height: 1, Workers: uint(workers)}
+		err = service.CreateSnapshot(params)
+		if err != nil {
+			t.Fatal(err)
+		}
 	}
-	fixture_path := filepath.Join(wd, "..", "..", "fixture")
-	datadir := filepath.Join(fixture_path, "chaindata")
-	if _, err := os.Stat(datadir); err != nil {
-		t.Fatal("no chaindata:", err)
+
+	testCases := []int{1, 4, 16, 32}
+	for _, tc := range testCases {
+		t.Run("case", func(t *testing.T) { runCase(t, tc) })
 	}
-	config := testConfig(datadir, filepath.Join(datadir, "ancient"))
-	fmt.Printf("config: %+v %+v\n", config.DB, config.Eth)
+}
+
+func failingPublishStateNode(_ *snapt.Node, _ string, _ *big.Int, _ snapt.Tx) error {
+	return errors.New("failingPublishStateNode")
+}
 
-	edb, err := NewLevelDB(config.Eth)
-	if err != nil {
-		t.Fatal(err)
-	}
-	workers := 8
-
-	pub := mock.NewMockPublisher(t)
-	pub.EXPECT().PublishHeader(gomock.Eq(fixt.PublishHeader))
-	pub.EXPECT().BeginTx().
-		Times(workers)
-	pub.EXPECT().PrepareTxForBatch(gomock.Any(), gomock.Any()).
-		Times(workers)
-	pub.EXPECT().PublishStateNode(gomock.Any(), mock.AnyOf(int64(0), int64(1)), gomock.Any()).
-		Times(workers)
-	// TODO: fixtures for storage node
-	// pub.EXPECT().PublishStorageNode(gomock.Eq(fixt.StorageNode), gomock.Eq(int64(0)), gomock.Any())
-	pub.EXPECT().CommitTx(gomock.Any()).
-		Times(workers)
-
-	service, err := NewSnapshotService(edb, pub)
-	if err != nil {
-		t.Fatal(err)
+func TestRecovery(t *testing.T) {
+	runCase := func(t *testing.T, workers int) {
+		pub, tx := makeMocks(t)
+		pub.EXPECT().PublishHeader(gomock.Any()).AnyTimes()
+		pub.EXPECT().BeginTx().Return(tx, nil).AnyTimes()
+		pub.EXPECT().PrepareTxForBatch(gomock.Any(), gomock.Any()).Return(tx, nil).AnyTimes()
+		pub.EXPECT().PublishStateNode(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+			Times(workers).
+			DoAndReturn(failingPublishStateNode)
+		tx.EXPECT().Commit().AnyTimes()
+
+		config := testConfig(fixt.ChaindataPath, fixt.AncientdataPath)
+		edb, err := NewLevelDB(config.Eth)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer edb.Close()
+
+		recovery := filepath.Join(t.TempDir(), "recover.csv")
+		service, err := NewSnapshotService(edb, pub, recovery)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		params := SnapshotParams{Height: 1, Workers: uint(workers)}
+		err = service.CreateSnapshot(params)
+		if err == nil {
+			t.Fatal("expected an error")
+		}
+
+		if _, err = os.Stat(recovery); err != nil {
+			t.Fatal("cannot stat recovery file:", err)
+		}
+
+		pub.EXPECT().PublishStateNode(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()
+		err = service.CreateSnapshot(params)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = os.Stat(recovery)
+		if err == nil {
+			t.Fatal("recovery file still present")
+		} else {
+			if !os.IsNotExist(err) {
+				t.Fatal(err)
+			}
+		}
 	}
 
-	params := SnapshotParams{Height: 1, Workers: uint(workers)}
-	err = service.CreateSnapshot(params)
-	if err != nil {
-		t.Fatal(err)
+	testCases := []int{1, 4, 32}
+	for _, tc := range testCases {
+		t.Run("case", func(t *testing.T) { runCase(t, tc) })
 	}
 
-	// err = service.CreateLatestSnapshot(0)
-	// if err != nil {
-	// 	t.Fatal(err)
-	// }
 }

pkg/snapshot/tracker.go

@@ -0,0 +1,173 @@
+package snapshot
+
+import (
+	"encoding/csv"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/trie"
+	log "github.com/sirupsen/logrus"
+
+	iter "github.com/vulcanize/go-eth-state-node-iterator"
+)
+
+type trackedIter struct {
+	trie.NodeIterator
+	tracker *iteratorTracker
+}
+
+func (it *trackedIter) Next(descend bool) bool {
+	ret := it.NodeIterator.Next(descend)
+	if !ret {
+		it.tracker.stopChan <- it
+	}
+	return ret
+}
+
+type iteratorTracker struct {
+	recoveryFile string
+
+	startChan chan *trackedIter
+	stopChan  chan *trackedIter
+	started   map[*trackedIter]struct{}
+	stopped   []*trackedIter
+
+	haltChan chan struct{}
+	done     chan struct{}
+}
+
+func newTracker(file string, buf int) iteratorTracker {
+	return iteratorTracker{
+		recoveryFile: file,
+		startChan:    make(chan *trackedIter, buf),
+		stopChan:     make(chan *trackedIter, buf),
+		started:      map[*trackedIter]struct{}{},
+		haltChan:     make(chan struct{}),
+		done:         make(chan struct{}),
+	}
+}
+
+func (tr *iteratorTracker) captureSignal() {
+	sigChan := make(chan os.Signal, 1)
+
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		sig := <-sigChan
+		log.Errorf("Signal received (%v), stopping", sig)
+		tr.haltAndDump()
+		os.Exit(1)
+	}()
+}
+
+// listens for starts/stops and manages current state
+func (tr *iteratorTracker) run() {
+loop:
+	for {
+		select {
+		case start := <-tr.startChan:
+			tr.started[start] = struct{}{}
+		case stop := <-tr.stopChan:
+			tr.stopped = append(tr.stopped, stop)
+		case <-tr.haltChan:
+			break loop
+		default:
+		}
+	}
+	tr.done <- struct{}{}
+}
+
+func (tr *iteratorTracker) tracked(it trie.NodeIterator) (ret *trackedIter) {
+	ret = &trackedIter{it, tr}
+	tr.startChan <- ret
+	return
+}
+
+// dumps iterator path and bounds to a text file so it can be restored later
+func (tr *iteratorTracker) dump() error {
+	log.Info("Dumping recovery state to:", tr.recoveryFile)
+	var rows [][]string
+	for it, _ := range tr.started {
+		var endPath []byte
+		if impl, ok := it.NodeIterator.(*iter.PrefixBoundIterator); ok {
+			endPath = impl.EndPath
+		}
+		rows = append(rows, []string{
+			fmt.Sprintf("%x", it.Path()),
+			fmt.Sprintf("%x", endPath),
+		})
+	}
+	file, err := os.Create(tr.recoveryFile)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+	out := csv.NewWriter(file)
+	return out.WriteAll(rows)
+}
+
+// attempts to read iterator state from file
+// if file doesn't exist, returns an empty slice with no error
+func (tr *iteratorTracker) restore(tree state.Trie) ([]trie.NodeIterator, error) {
+	file, err := os.Open(tr.recoveryFile)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil, nil
+		}
+		return nil, err
+	}
+	defer file.Close()
+	in := csv.NewReader(file)
+	in.FieldsPerRecord = 2
+	rows, err := in.ReadAll()
+	if err != nil {
+		return nil, err
+	}
+	var ret []trie.NodeIterator
+	for _, row := range rows {
+		// pick up where each interval left off
+		var paths [2][]byte
+		for i, val := range row {
+			if len(val) != 0 {
+				if _, err = fmt.Sscanf(val, "%x", &paths[i]); err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		it := iter.NewPrefixBoundIterator(tree, paths[0], paths[1])
+		ret = append(ret, tr.tracked(it))
+	}
+	return ret, nil
+}
+
+func (tr *iteratorTracker) haltAndDump() error {
+	tr.haltChan <- struct{}{}
+	<-tr.done
+
+	// drain any pending events
+	close(tr.startChan)
+	for start := range tr.startChan {
+		tr.started[start] = struct{}{}
+	}
+	close(tr.stopChan)
+	for stop := range tr.stopChan {
+		tr.stopped = append(tr.stopped, stop)
+	}
+
+	for _, stop := range tr.stopped {
+		delete(tr.started, stop)
+	}
+
+	if len(tr.started) == 0 {
+		// if the tracker state is empty, erase any existing recovery file
+		err := os.Remove(tr.recoveryFile)
+		if os.IsNotExist(err) {
+			err = nil
+		}
+		return err
+	}
+	return tr.dump()
+}

pkg/snapshot/util.go

@@ -0,0 +1,26 @@
+package snapshot
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
+
+	file "github.com/vulcanize/ipld-eth-state-snapshot/pkg/snapshot/file"
+	pg "github.com/vulcanize/ipld-eth-state-snapshot/pkg/snapshot/pg"
+	snapt "github.com/vulcanize/ipld-eth-state-snapshot/pkg/types"
+)
+
+func NewPublisher(mode SnapshotMode, config *Config) (snapt.Publisher, error) {
+	switch mode {
+	case PgSnapshot:
+		driver, err := postgres.NewPGXDriver(context.Background(), config.DB.ConnConfig, config.Eth.NodeInfo)
+		if err != nil {
+			return nil, err
+		}
+		return pg.NewPublisher(postgres.NewPostgresDB(driver)), nil
+	case FileSnapshot:
+		return file.NewPublisher(config.File.OutputDir, config.Eth.NodeInfo)
+	}
+	return nil, fmt.Errorf("invalid snapshot mode: %s", mode)
+}

pkg/types/node_type.go

@@ -22,13 +22,12 @@ import (
 )
 
 // node for holding trie node information
-type node struct {
+type Node struct {
 	NodeType nodeType
 	Path     []byte
 	Key      common.Hash
 	Value    []byte
 }
-type Node = node
 
 // nodeType for explicitly setting type of node
 type nodeType int

pkg/types/publisher.go

@@ -1,17 +1,22 @@
 package types
 
 import (
+	"math/big"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
-	"github.com/jmoiron/sqlx"
 )
 
 type Publisher interface {
-	PublishHeader(header *types.Header) (int64, error)
-	PublishStateNode(node *Node, headerID int64, tx *sqlx.Tx) (int64, error)
-	PublishStorageNode(node *Node, stateID int64, tx *sqlx.Tx) error
-	PublishCode(codeHash common.Hash, codeBytes []byte, tx *sqlx.Tx) error
-	BeginTx() (*sqlx.Tx, error)
-	CommitTx(*sqlx.Tx) error
-	PrepareTxForBatch(tx *sqlx.Tx, batchSize uint) (*sqlx.Tx, error)
+	PublishHeader(header *types.Header) error
+	PublishStateNode(node *Node, headerID string, height *big.Int, tx Tx) error
+	PublishStorageNode(node *Node, headerID string, height *big.Int, statePath []byte, tx Tx) error
+	PublishCode(height *big.Int, codeHash common.Hash, codeBytes []byte, tx Tx) error
+	BeginTx() (Tx, error)
+	PrepareTxForBatch(tx Tx, batchSize uint) (Tx, error)
+}
+
+type Tx interface {
+	Rollback() error
+	Commit() error
 }

pkg/types/schema.go

@@ -0,0 +1,76 @@
+package types
+
+var TableIPLDBlock = Table{
+	`public.blocks`,
+	[]column{
+		{"block_number", bigint},
+		{"key", text},
+		{"data", bytea},
+	},
+	"ON CONFLICT (key, block_number) DO NOTHING",
+}
+
+var TableNodeInfo = Table{
+	Name: `public.nodes`,
+	Columns: []column{
+		{"genesis_block", varchar},
+		{"network_id", varchar},
+		{"node_id", varchar},
+		{"client_name", varchar},
+		{"chain_id", integer},
+	},
+}
+
+var TableHeader = Table{
+	"eth.header_cids",
+	[]column{
+		{"block_number", bigint},
+		{"block_hash", varchar},
+		{"parent_hash", varchar},
+		{"cid", text},
+		{"td", numeric},
+		{"node_id", varchar},
+		{"reward", numeric},
+		{"state_root", varchar},
+		{"tx_root", varchar},
+		{"receipt_root", varchar},
+		{"uncle_root", varchar},
+		{"bloom", bytea},
+		{"timestamp", numeric},
+		{"mh_key", text},
+		{"times_validated", integer},
+		{"coinbase", varchar},
+	},
+	"ON CONFLICT (block_hash, block_number) DO UPDATE SET (parent_hash, cid, td, node_id, reward, state_root, tx_root, receipt_root, uncle_root, bloom, timestamp, mh_key, times_validated, coinbase) = (EXCLUDED.parent_hash, EXCLUDED.cid, EXCLUDED.td, EXCLUDED.node_id, EXCLUDED.reward, EXCLUDED.state_root, EXCLUDED.tx_root, EXCLUDED.receipt_root, EXCLUDED.uncle_root, EXCLUDED.bloom, EXCLUDED.timestamp, EXCLUDED.mh_key, eth.header_cids.times_validated + 1, EXCLUDED.coinbase)",
+}
+
+var TableStateNode = Table{
+	"eth.state_cids",
+	[]column{
+		{"block_number", bigint},
+		{"header_id", varchar},
+		{"state_leaf_key", varchar},
+		{"cid", text},
+		{"state_path", bytea},
+		{"node_type", integer},
+		{"diff", boolean},
+		{"mh_key", text},
+	},
+	"ON CONFLICT (header_id, state_path, block_number) DO UPDATE SET (state_leaf_key, cid, node_type, diff, mh_key) = (EXCLUDED.state_leaf_key, EXCLUDED.cid, EXCLUDED.node_type, EXCLUDED.diff, EXCLUDED.mh_key)",
+}
+
+var TableStorageNode = Table{
+	"eth.storage_cids",
+	[]column{
+		{"block_number", bigint},
+		{"header_id", varchar},
+		{"state_path", bytea},
+		{"storage_leaf_key", varchar},
+		{"cid", text},
+		{"storage_path", bytea},
+		{"node_type", integer},
+		{"diff", boolean},
+		{"mh_key", text},
+	},
+	"ON CONFLICT (header_id, state_path, storage_path, block_number) DO UPDATE SET (storage_leaf_key, cid, node_type, diff, mh_key) = (EXCLUDED.storage_leaf_key, EXCLUDED.cid, EXCLUDED.node_type, EXCLUDED.diff, EXCLUDED.mh_key)",
+}

pkg/types/table.go

@@ -0,0 +1,79 @@
+package types
+
+import (
+	"fmt"
+	"strings"
+)
+
+type colType int
+
+const (
+	integer colType = iota
+	boolean
+	bigint
+	numeric
+	bytea
+	varchar
+	text
+)
+
+type column struct {
+	name string
+	typ  colType
+}
+type Table struct {
+	Name           string
+	Columns        []column
+	conflictClause string
+}
+
+func (tbl *Table) ToCsvRow(args ...interface{}) []string {
+	var row []string
+	for i, col := range tbl.Columns {
+		row = append(row, col.typ.formatter()(args[i]))
+	}
+	return row
+}
+
+func (tbl *Table) ToInsertStatement() string {
+	var colnames, placeholders []string
+	for i, col := range tbl.Columns {
+		colnames = append(colnames, col.name)
+		placeholders = append(placeholders, fmt.Sprintf("$%d", i+1))
+	}
+	return fmt.Sprintf(
+		"INSERT INTO %s (%s) VALUES (%s) %s",
+		tbl.Name, strings.Join(colnames, ", "), strings.Join(placeholders, ", "), tbl.conflictClause,
+	)
+}
+
+type colfmt = func(interface{}) string
+
+func sprintf(f string) colfmt {
+	return func(x interface{}) string { return fmt.Sprintf(f, x) }
+}
+
+func (typ colType) formatter() colfmt {
+	switch typ {
+	case integer:
+		return sprintf("%d")
+	case boolean:
+		return func(x interface{}) string {
+			if x.(bool) {
+				return "t"
+			}
+			return "f"
+		}
+	case bigint:
+		return sprintf("%s")
+	case numeric:
+		return sprintf("%d")
+	case bytea:
+		return sprintf(`\x%x`)
+	case varchar:
+		return sprintf("%s")
+	case text:
+		return sprintf("%s")
+	}
+	panic("unreachable")
+}

pkg/types/util.go

@@ -0,0 +1,33 @@
+package types
+
+import (
+	"bytes"
+
+	"github.com/sirupsen/logrus"
+
+	"github.com/ethereum/go-ethereum/common"
+)
+
+var nullHash = common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000000")
+
+func IsNullHash(hash common.Hash) bool {
+	return bytes.Equal(hash.Bytes(), nullHash.Bytes())
+}
+
+func CommitOrRollback(tx Tx, err error) error {
+	var rberr error
+	defer func() {
+		if rberr != nil {
+			logrus.Errorf("rollback failed: %s", rberr)
+		}
+	}()
+	if rec := recover(); rec != nil {
+		rberr = tx.Rollback()
+		panic(rec)
+	} else if err != nil {
+		rberr = tx.Rollback()
+	} else {
+		err = tx.Commit()
+	}
+	return err
+}

test/helper.go

@@ -0,0 +1,59 @@
+package test
+
+import (
+	"bytes"
+	"os"
+	"reflect"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
+	ethnode "github.com/ethereum/go-ethereum/statediff/indexer/node"
+)
+
+var (
+	DefaultNodeInfo = ethnode.Info{
+		ID:           "test_nodeid",
+		ClientName:   "test_client",
+		GenesisBlock: "TEST_GENESIS",
+		NetworkID:    "test_network",
+		ChainID:      0,
+	}
+	DefaultPgConfig = postgres.Config{
+		Hostname:     "localhost",
+		Port:         5432,
+		DatabaseName: "vulcanize_test",
+		Username:     "vulcanize",
+		Password:     "vulcanize_password",
+
+		MaxIdle:         0,
+		MaxConnLifetime: 0,
+		MaxConns:        4,
+	}
+)
+
+func NeedsDB(t *testing.T) {
+	t.Helper()
+	if os.Getenv("TEST_WITH_DB") == "" {
+		t.Skip("set TEST_WITH_DB to enable test")
+	}
+}
+
+func NoError(t *testing.T, err error) {
+	t.Helper()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+// ExpectEqual asserts the provided interfaces are deep equal
+func ExpectEqual(t *testing.T, want, got interface{}) {
+	if !reflect.DeepEqual(want, got) {
+		t.Fatalf("Values not equal:\nExpected:\t%v\nActual:\t\t%v", want, got)
+	}
+}
+
+func ExpectEqualBytes(t *testing.T, want, got []byte) {
+	if !bytes.Equal(want, got) {
+		t.Fatalf("Bytes not equal:\nExpected:\t%v\nActual:\t\t%v", want, got)
+	}
+}