wip

chiv.go

@@ -1,4 +1,4 @@
-// Package chiv archives arbitrarily large relational database tables to Amazon S3.
+// Package chiv archives relational database tables to Amazon S3.
 package chiv
 
 import (
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/s3/s3manager"
@@ -19,8 +20,20 @@ var (
 	ErrRecordLength = errors.New("record length does not match number of columns")
 	// ErrParserRegex initialization problem.
 	ErrParserRegex = errors.New("initializing parser regex")
+	// ErrBuildingQuery string.
+	ErrBuildingQuery = errors.New("building query")
 )
 
+// Archive a database table to S3.
+func Archive(db *sql.DB, s3 *s3manager.Uploader, table, bucket string, options ...Option) error {
+	return NewArchiver(db, s3).ArchiveWithContext(context.Background(), table, bucket, options...)
+}
+
+// ArchiveWithContext is like Archive, with context.
+func ArchiveWithContext(ctx context.Context, db *sql.DB, s3 *s3manager.Uploader, table, bucket string, options ...Option) error {
+	return NewArchiver(db, s3).ArchiveWithContext(ctx, table, bucket, options...)
+}
+
 // Archiver archives arbitrarily large relational database tables to Amazon S3.
 type Archiver struct {
 	db        *sql.DB
@@ -29,6 +42,7 @@ type Archiver struct {
 	key       string
 	extension string
 	null      []byte
+	columns   []string
 }
 
 // NewArchiver constructs an archiver with the given database, S3 uploader and options.
@@ -47,7 +61,7 @@ func NewArchiver(db *sql.DB, s3 *s3manager.Uploader, options ...Option) *Archive
 	return &a
 }
 
-// Archive a database table to S3.
+// Archive a database table to S3. Any options provided override those set on creation.
 func (a *Archiver) Archive(table, bucket string, options ...Option) error {
 	return a.ArchiveWithContext(context.Background(), table, bucket, options...)
 }
@@ -80,8 +94,7 @@ func (a *Archiver) archive(ctx context.Context, table string, bucket string) err
 }
 
 func (a *Archiver) download(ctx context.Context, wc io.WriteCloser, table string, errs chan error) {
-	selectAll := fmt.Sprintf(`select * from "%s";`, table)
-	rows, err := a.db.QueryContext(ctx, selectAll)
+	rows, err := a.query(ctx, table)
 	if err != nil {
 		errs <- err
 		return
@@ -146,6 +159,24 @@ func (a *Archiver) download(ctx context.Context, wc io.WriteCloser, table string
 	}
 }
 
+func (a *Archiver) query(ctx context.Context, table string) (*sql.Rows, error) {
+	var b strings.Builder
+	for i, column := range a.columns {
+		b.WriteString(column)
+		if i < len(a.columns)-1 {
+			b.WriteString(", ")
+		}
+	}
+
+	columns := "*"
+	if b.Len() > 0 {
+		columns = b.String()
+	}
+
+	query := fmt.Sprintf(`select %s from "%s";`, columns, table)
+	return a.db.QueryContext(ctx, query)
+}
+
 func (a *Archiver) upload(ctx context.Context, r io.Reader, table string, bucket string, errs chan error) {
 	if a.key == "" {
 		if a.extension != "" {

chiv_benchmark_test.go

@@ -31,7 +31,7 @@ func BenchmarkArchiver_Archive(b *testing.B) {
 	)
 
 	var (
-		benchmarks = []int{1, 10, 100, 1000, 3000, 5000, 10000}
+		benchmarks = []int{1, 10, 100, 1000, 3000}
 
 		db       = newDB(b, "postgres", os.Getenv("POSTGRES_URL"))
 		s3client = newS3Client(b, os.Getenv("AWS_REGION"), os.Getenv("AWS_ENDPOINT"))
@@ -48,12 +48,10 @@ func BenchmarkArchiver_Archive(b *testing.B) {
 			exec(b, db, statement)
 		}
 
-		subject := chiv.NewArchiver(db, uploader)
-
 		b.Run(fmt.Sprintf("benchmark_%d", count), func(*testing.B) {
 			for j := 0; j < b.N; j++ {
 				key := fmt.Sprintf("benchmark_%d_%d", count, j)
-				if err := subject.Archive(table, bucket, chiv.WithKey(key)); err != nil {
+				if err := chiv.Archive(db, uploader, table, bucket, chiv.WithKey(key)); err != nil {
 					b.Error(err)
 				}
 			}

chiv_test.go

@@ -3,6 +3,7 @@
 package chiv_test
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -35,6 +36,23 @@ type call struct {
 
 func TestArchiver_Archive(t *testing.T) {
 	cases := []test{
+		{
+			name:     "happy path csv",
+			driver:   "postgres",
+			database: os.Getenv("POSTGRES_URL"),
+			setup:    "./test/data/database_setup.sql",
+			teardown: "./test/data/database_teardown.sql",
+			bucket:   "database_bucket",
+			options:  []chiv.Option{},
+			calls: []call{
+				{
+					expected: "./test/data/database.csv",
+					table:    "database_table",
+					key:      "database_table",
+					options:  []chiv.Option{},
+				},
+			},
+		},
 		{
 			name:     "postgres to csv",
 			driver:   "postgres",
@@ -160,23 +178,78 @@ func TestArchiver_Archive(t *testing.T) {
 				},
 			},
 		},
-		//{
-		//	name:     "mysql to csv",
-		//	driver:   "mysql",
-		//	database: os.Getenv("MYSQL_URL"),
-		//	setup:    "./test/data/mysql_setup.sql",
-		//	teardown: "./test/data/mysql_teardown.sql",
-		//	bucket:   "mysql_bucket",
-		//	options:  []chiv.Option{},
-		//	calls: []call{
-		//		{
-		//			expected: "./test/data/mysql.csv",
-		//			table:    "mysql_table",
-		//			key:      "mysql_table",
-		//			options:  []chiv.Option{},
-		//		},
-		//	},
-		//},
+		{
+			name:     "postgres one-off extension",
+			driver:   "postgres",
+			database: os.Getenv("POSTGRES_URL"),
+			setup:    "./test/data/postgres_setup.sql",
+			teardown: "./test/data/postgres_teardown.sql",
+			bucket:   "postgres_bucket",
+			options: []chiv.Option{
+				chiv.WithFormat(chiv.YAML),
+			},
+			calls: []call{
+				{
+					expected: "./test/data/postgres.yaml",
+					table:    "postgres_table",
+					key:      "postgres_table.not_yaml",
+					options: []chiv.Option{
+						chiv.WithExtension("not_yaml"),
+					},
+				},
+				{
+					expected: "./test/data/postgres.yaml",
+					table:    "postgres_table",
+					key:      "postgres_table",
+					options:  []chiv.Option{},
+				},
+			},
+		},
+		{
+			name:     "postgres two tables",
+			driver:   "postgres",
+			database: os.Getenv("POSTGRES_URL"),
+			setup:    "./test/data/two_tables_setup.sql",
+			teardown: "./test/data/two_tables_teardown.sql",
+			bucket:   "postgres_bucket",
+			options: []chiv.Option{
+				chiv.WithFormat(chiv.CSV),
+				chiv.WithExtension("csv"),
+			},
+			calls: []call{
+				{
+					expected: "./test/data/two_tables_first.csv",
+					table:    "first_table",
+					key:      "first_table.csv",
+					options:  []chiv.Option{},
+				},
+				{
+					expected: "./test/data/two_tables_second.csv",
+					table:    "second_table",
+					key:      "second_table.csv",
+					options:  []chiv.Option{},
+				},
+			},
+		},
+		{
+			name:     "with columns",
+			driver:   "postgres",
+			database: os.Getenv("POSTGRES_URL"),
+			setup:    "./test/data/postgres_setup.sql",
+			teardown: "./test/data/postgres_teardown.sql",
+			bucket:   "postgres_bucket",
+			options:  []chiv.Option{},
+			calls: []call{
+				{
+					expected: "./test/data/postgres_subset.csv",
+					table:    "postgres_table",
+					key:      "postgres_table",
+					options: []chiv.Option{
+						chiv.WithColumns("id", "text_column", "int_column"),
+					},
+				},
+			},
+		},
 	}
 
 	for _, test := range cases {
@@ -207,3 +280,30 @@ func TestArchiver_Archive(t *testing.T) {
 		})
 	}
 }
+
+func TestArchiveWithContext(t *testing.T) {
+	var (
+		database   = os.Getenv("POSTGRES_URL")
+		driver     = "postgres"
+		bucket     = "postgres_bucket"
+		table      = "postgres_table"
+		setup      = "./test/data/postgres_setup.sql"
+		teardown   = "./test/data/postgres_teardown.sql"
+		expected   = "./test/data/postgres.csv"
+		db         = newDB(t, driver, database)
+		s3client   = newS3Client(t, os.Getenv("AWS_REGION"), os.Getenv("AWS_ENDPOINT"))
+		uploader   = s3manager.NewUploaderWithClient(s3client)
+		downloader = s3manager.NewDownloaderWithClient(s3client)
+	)
+
+	exec(t, db, readFile(t, setup))
+	defer exec(t, db, readFile(t, teardown))
+
+	createBucket(t, s3client, bucket)
+	defer deleteBucket(t, s3client, bucket)
+
+	require.NoError(t, chiv.ArchiveWithContext(context.Background(), db, uploader, table, bucket))
+
+	actual := download(t, downloader, bucket, table)
+	require.Equal(t, readFile(t, expected), actual)
+}

cmd/chiv/main.go

@@ -1,4 +1,5 @@
 package main
 
 func main() {
+
 }

codeship-services.yml

@@ -1,5 +1,7 @@
 test:
-  build: .
+  build:
+    context: .
+    dockerfile: ./test/Dockerfile
   cached: true
   environment:
     - AWS_ACCESS_KEY_ID=bogus
@@ -12,8 +14,11 @@ test:
     - s3
 
 lint:
-  build: .
+  build:
+    context: .
+    dockerfile: ./test/Dockerfile
   cached: true
+  command: golangci-lint run
 
 postgres:
   image: healthcheck/postgres

go.mod

@@ -3,13 +3,13 @@ module github.com/gavincabbage/chiv
 go 1.12
 
 require (
+	github.com/Masterminds/squirrel v1.1.0
 	github.com/aws/aws-sdk-go v1.19.6
-	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/go-sql-driver/mysql v1.4.1
+	github.com/golangci/golangci-lint v1.16.0
 	github.com/lib/pq v1.0.0
+	github.com/mattn/go-sqlite3 v1.10.0 // indirect
 	github.com/stretchr/testify v1.3.0
-	golang.org/x/net v0.0.0-20190313220215-9f648a60d977 // indirect
-	golang.org/x/tour v0.0.0-20190318020441-db40fe78fefc // indirect
 	google.golang.org/appengine v1.5.0 // indirect
 	gopkg.in/yaml.v2 v2.2.2
 )