bulk-sql-insert: just some quick script to bulk insert records into a sql database with some throttling and batching capabilities

main.go

package main

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"time"

	"gopkg.in/cheggaaa/pb.v1"

	_ "github.com/denisenkom/go-mssqldb"
)

var (
	dbURI  string
	useTx  bool
	txSize int
	db     *sql.DB
)

func init() {
	flag.StringVar(&dbURI, "database", "", "The database URI to insert records to")
	flag.BoolVar(&useTx, "transaction", false, "Use transaction style insert")
	flag.IntVar(&txSize, "bucketSize", -1, "Max size of records per transactions")
}

func newErr(msg string) error {
	return fmt.Errorf("Error: %s\n\nUsage: %s --database=... /path/to/insert_file.sql", msg, os.Args[0])
}

func main() {
	flag.Parse()
	if flag.NArg() < 1 {
		log.Fatal(newErr("insert file required"))
	}

	if len(dbURI) == 0 {
		log.Fatal(newErr("database URI required"))
	}

	err := connect(dbURI)
	if err != nil {
		log.Fatal(err)
	}

	now := time.Now()

	fmt.Printf("Opening %s...\n", flag.Arg(0))
	file, err := openFile(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	var (
		lines    = strings.Split(file.String(), "\n")
		bar      = pb.New(len(lines))
		totalNum = 0
		i        = 0
	)
	bar.Start()
	for _, line := range lines {
		res, err := db.Exec(line)
		if err != nil {
			log.Fatal(err)
		}
		n, err := res.RowsAffected()
		if err != nil {
			log.Fatal(err)
		}
		totalNum += int(n)
		bar.Increment()
		i++
	}
	bar.FinishPrint(fmt.Sprintf("%d records inserted from %d statements in %s\n", totalNum, i, time.Since(now)))
}

func connect(uri string) (err error) {
	db, err = sql.Open("sqlserver", uri)
	return
}

func openFile(filePath string) (*bytes.Buffer, error) {
	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}

	return bytes.NewBuffer(data), err
}