Example of copying a CSV file into Postgres

pq_copy_demo.go

package main

import (
	"database/sql"
	"encoding/csv"
	"io"
	"log"
	"os"
	"strconv"

	"github.com/lib/pq"
	_ "github.com/lib/pq"
)

// A standalone test that
// - reads a specific csv file with a known schema
// - uses the `pq` driver to connect to a local database
// - utilizes `COPY ...` to pass through the rows read from the csv file
// - reads back the rows read to the table

var (
	csvFileName = "/Users/agam/tmp/upload_test_noheader.csv"
)

func main() {
	// Connect.
	connStr := "postgres://postgres:mypostgres@localhost:5433/postgres?sslmode=disable"
	pgDB, err := sql.Open("postgres", connStr)
	if err != nil {
		log.Fatal(err)
	}

	// Basic SQL sanity check.
	rows, err := pgDB.Query("SELECT * FROM foo")
	if err != nil {
		log.Fatal(err)
	}

	rowCount := 0
	for rows.Next() {
		rowCount++
	}
	log.Printf("Debug: got %d rows\n", rowCount)

	// Get our CSV data.
	csvFile, err := os.Open(csvFileName)
	if err != nil {
		log.Fatal(err)
	}
	csvReader := csv.NewReader(csvFile)

	// Construct our Copy statement.
	copyStmt := pq.CopyIn("upload_test", "foo", "bar")
	log.Printf("Copy stmt: %s\n", copyStmt)

	// Prepare to execute.
	tx, err := pgDB.Begin()
	if err != nil {
		log.Fatal(err)
	}
	copyTx, err := tx.Prepare(copyStmt)
	if err != nil {
		log.Fatal(err)
	}

	// "Execute" by iterating through records.
	for {
		// Read each record from csv
		record, err := csvReader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		// "Pump" this record in to our copy-stmt.
		// Note: this implies _knowing the schema_ (which we would for Avro!)
		fooVal, _ := strconv.Atoi(record[0])
		barVal, _ := strconv.Atoi(record[1])
		_, err = copyTx.Exec(fooVal, barVal)
		if err != nil {
			log.Fatal(err)
		}
	}

	result, err := copyTx.Exec()
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("Debug: exec-result: %v\n", result)

	// We're done!
	err = copyTx.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal(err)
	}

	// Read back what we wrote.
	rows, err = pgDB.Query("SELECT * FROM upload_test")
	if err != nil {
		log.Fatal(err)
	}

	rowCount = 0
	for rows.Next() {
		rowCount++
	}
	log.Printf("Debug: successfully wrote %d rows\n", rowCount)
}

// Sample run:
// ~/code/golang/pqcopy master !1 ❯ go run pqcopy.go                                                                                                                                                                    01:15:05
// 2021/10/27 01:16:33 Debug: got 4 rows
// 2021/10/27 01:16:33 Copy stmt: COPY "upload_test" ("foo", "bar") FROM STDIN
// 2021/10/27 01:16:33 Debug: exec-result: {0xc0000ce000 3}
// 2021/10/27 01:16:33 Debug: successfully wrote 3 rows

// Sample file used:
// 2,3
// 10,20
// 4,5