csv2pg.go

package main

import (
	"database/sql"
	"encoding/csv"
	"fmt"
	"github.com/lib/pq"
	"io"
	"log"
	"os"
	"strings"
	"time"
)

func main() {
	if len(os.Args) < 3 {
		log.Fatalf("Usage: %s <file> <table_name>", os.Args[0])
	}
	fobj, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	tableName := os.Args[2]
	db, err := sql.Open("postgres", "")
	if err != nil {
		log.Fatal("Cannot open database: ", err)
	}

	startTime := time.Now()
	rowCount := 0

	csvReader := csv.NewReader(fobj)
	var copyTx *sql.Tx = nil
	var copyStmt *sql.Stmt = nil
	var csvHeader []string = nil
	for {
		rec, err := csvReader.Read()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatal("Error reading: ", err)
		}
		if csvHeader == nil {
			csvHeader = rec
			createCols := []string{}
			for _, col := range csvHeader {
				createCols = append(createCols, fmt.Sprintf(`"%s" text`, col))
			}
			log.Printf("Using table %#v with cols: %s", tableName, csvHeader)
			_, err := db.Exec(fmt.Sprintf(`CREATE TABLE "%s" (
				id bigserial primary key,
				stamp timestamptz default now(),
				%s
				)`, tableName, strings.Join(createCols, ", ")))
			if err != nil {
				if pgerr, ok := err.(*pq.Error); ok && pgerr.Code.Name() == "duplicate_table" {
					continue
				}
				log.Fatal("Could not create table: ", err)
			}
			continue
		}
		if copyTx == nil {
			copyTx, err = db.Begin()
			if err != nil {
				log.Fatal("Cannot start transaction: ", err)
			}
			copyStmt, err = copyTx.Prepare(pq.CopyIn(tableName, csvHeader...))
			if err != nil {
				log.Fatal("Error creating CopyIn statement: ", err)
			}
		}
		copyRec := make([]interface{}, len(rec))
		for i := range rec {
			copyRec[i] = rec[i]
		}
		if _, err := copyStmt.Exec(copyRec...); err != nil {
			log.Fatalf("CopyIn failed (row=%d): %s", rowCount, err)
		}
		rowCount++
	}
	if copyTx != nil {
		if _, err := copyStmt.Exec(); err != nil {
			log.Fatal("CopyIn Exec failed: ", err)
		}
		if err := copyStmt.Close(); err != nil {
			log.Fatal("CopyIn Close failed: ", err)
		}
		if err := copyTx.Commit(); err != nil {
			log.Fatal("Transaction commit failed: ", err)
		}
	}
	log.Printf("Copied in %d records in %s", rowCount, time.Now().Sub(startTime))
}