Add helper package to parse multi-statement migrations

dhui · dhui · commit 620492183fbf · 2020-07-07T22:50:09.000-07:00
Addresses: golang-migrate#406
diff --git a/database/multistmt/parse.go b/database/multistmt/parse.go
@@ -0,0 +1,46 @@
+// Package multistmt provides methods for parsing multi-statement database migrations
+package multistmt
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+)
+
+// StartBufSize is the default starting size of the buffer used to scan and parse multi-statement migrations
+var StartBufSize = 4096
+
+// Handler handles a single migration parsed from a multi-statement migration.
+// It's given the single migration to handle and returns whether or not further statements
+// from the multi-statement migration should be parsed and handled.
+type Handler func(migration []byte) bool
+
+func splitWithDelimiter(delimiter []byte) func(d []byte, atEOF bool) (int, []byte, error) {
+	return func(d []byte, atEOF bool) (int, []byte, error) {
+		// SplitFunc inspired by bufio.ScanLines() implementation
+		if atEOF {
+			if len(d) == 0 {
+				return 0, nil, nil
+			}
+			return len(d), d, nil
+		}
+		if i := bytes.Index(d, delimiter); i >= 0 {
+			return i + len(delimiter), d[:i+len(delimiter)], nil
+		}
+		return 0, nil, nil
+	}
+}
+
+// Parse parses the given multi-statement migration
+func Parse(reader io.Reader, delimiter []byte, maxMigrationSize int, h Handler) error {
+	scanner := bufio.NewScanner(reader)
+	scanner.Buffer(make([]byte, 0, StartBufSize), maxMigrationSize)
+	scanner.Split(splitWithDelimiter(delimiter))
+	for scanner.Scan() {
+		cont := h(scanner.Bytes())
+		if !cont {
+			break
+		}
+	}
+	return scanner.Err()
+}
diff --git a/database/multistmt/parse_test.go b/database/multistmt/parse_test.go
@@ -0,0 +1,57 @@
+package multistmt_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/golang-migrate/migrate/v4/database/multistmt"
+)
+
+const maxMigrationSize = 1024
+
+func TestParse(t *testing.T) {
+	testCases := []struct {
+		name        string
+		multiStmt   string
+		delimiter   string
+		expected    []string
+		expectedErr error
+	}{
+		{name: "single statement, no delimiter", multiStmt: "single statement, no delimiter", delimiter: ";",
+			expected: []string{"single statement, no delimiter"}, expectedErr: nil},
+		{name: "single statement, one delimiter", multiStmt: "single statement, one delimiter;", delimiter: ";",
+			expected: []string{"single statement, one delimiter;"}, expectedErr: nil},
+		{name: "two statements, no trailing delimiter", multiStmt: "statement one; statement two", delimiter: ";",
+			expected: []string{"statement one;", " statement two"}, expectedErr: nil},
+		{name: "two statements, with trailing delimiter", multiStmt: "statement one; statement two;", delimiter: ";",
+			expected: []string{"statement one;", " statement two;"}, expectedErr: nil},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			stmts := make([]string, 0, len(tc.expected))
+			err := multistmt.Parse(strings.NewReader(tc.multiStmt), []byte(tc.delimiter), maxMigrationSize, func(b []byte) bool {
+				stmts = append(stmts, string(b))
+				return true
+			})
+			assert.Equal(t, tc.expectedErr, err)
+			assert.Equal(t, tc.expected, stmts)
+		})
+	}
+}
+
+func TestParseDiscontinue(t *testing.T) {
+	multiStmt := "statement one; statement two"
+	delimiter := ";"
+	expected := []string{"statement one;"}
+
+	stmts := make([]string, 0, len(expected))
+	err := multistmt.Parse(strings.NewReader(multiStmt), []byte(delimiter), maxMigrationSize, func(b []byte) bool {
+		stmts = append(stmts, string(b))
+		return false
+	})
+	assert.Nil(t, err)
+	assert.Equal(t, expected, stmts)
+}
