Merge pull request #2 from arthurkushman/timestamp_migrations_dmls

#1: Add an ability to call INSERT stmts in spanner db

pkg/spanner/client.go

@@ -27,7 +27,7 @@ import (
 
 	"cloud.google.com/go/spanner"
 	databasev1 "cloud.google.com/go/spanner/admin/database/apiv1"
-	databasepb "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
+	"cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
 	"github.com/hashicorp/go-multierror"
 	"google.golang.org/api/iterator"
@@ -356,6 +356,13 @@ func (c *Client) ExecuteMigrations(ctx context.Context, migrations Migrations, l
 				}
 			}
 		case statementKindDML:
+			if _, err := c.ApplyDML(ctx, m.Statements, PriorityTypeHigh); err != nil {
+				return &Error{
+					Code: ErrorCodeExecuteMigrations,
+					err:  err,
+				}
+			}
+		case statementKindPartitionedDML:
 			if _, err := c.ApplyPartitionedDML(ctx, m.Statements, PriorityTypeUnspecified); err != nil {
 				return &Error{
 					Code: ErrorCodeExecuteMigrations,

pkg/spanner/migration.go

@@ -22,7 +22,7 @@ package spanner
 import (
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"os"
 	"path/filepath"
 	"regexp"
 	"strconv"
@@ -39,12 +39,18 @@ var (
 
 	MigrationNameRegex = regexp.MustCompile(`[a-zA-Z0-9_\-]+`)
 
-	dmlRegex = regexp.MustCompile("^(UPDATE|DELETE)[\t\n\f\r ].*")
+	dmlRegex = regexp.MustCompile("^(INSERT|UPDATE|DELETE)[\t\n\f\r ].*")
+
+	// Instead of a regex to determine if a DML statement is a PartitionedDML we check if it's
+	// not a PartitionedDML by checking for INSERT statements.
+	// An improvement would be to use spanners algo to distinguish between DML types.
+	notPartitionedDmlRegex = regexp.MustCompile(`(?i)^INSERT`)
 )
 
 const (
-	statementKindDDL statementKind = "DDL"
-	statementKindDML statementKind = "DML"
+	statementKindDDL            statementKind = "DDL"
+	statementKindDML            statementKind = "DML"
+	statementKindPartitionedDML statementKind = "PartitionedDML"
 )
 
 type (
@@ -80,7 +86,7 @@ func (ms Migrations) Less(i, j int) bool {
 }
 
 func LoadMigrations(dir string) (Migrations, error) {
-	files, err := ioutil.ReadDir(dir)
+	files, err := os.ReadDir(dir)
 	if err != nil {
 		return nil, err
 	}
@@ -106,7 +112,7 @@ func LoadMigrations(dir string) (Migrations, error) {
 
 		fileName := f.Name()
 
-		file, err := ioutil.ReadFile(filepath.Join(dir, fileName))
+		file, err := os.ReadFile(filepath.Join(dir, fileName))
 		if err != nil {
 			continue
 		}
@@ -171,29 +177,61 @@ func dmlToStatements(filename string, data []byte) ([]string, error) {
 
 func inspectStatementsKind(statements []string) (statementKind, error) {
 	kindMap := map[statementKind]uint64{
-		statementKindDDL: 0,
-		statementKindDML: 0,
+		statementKindDDL:            0,
+		statementKindDML:            0,
+		statementKindPartitionedDML: 0,
 	}
 
 	for _, s := range statements {
-		if isDML(s) {
-			kindMap[statementKindDML]++
-		} else {
-			kindMap[statementKindDDL]++
-		}
+		kindMap[getStatementKind(s)]++
 	}
 
-	if kindMap[statementKindDML] > 0 {
-		if kindMap[statementKindDDL] > 0 {
-			return "", errors.New("cannot specify DDL and DML at same migration file")
-		}
+	if distinctKind(kindMap, statementKindDDL) {
+		return statementKindDDL, nil
+	}
 
+	if distinctKind(kindMap, statementKindDML) {
 		return statementKindDML, nil
 	}
 
-	return statementKindDDL, nil
+	if distinctKind(kindMap, statementKindPartitionedDML) {
+		return statementKindPartitionedDML, nil
+	}
+
+	return "", errors.New("cannot specify DDL and DML in the same migration file")
 }
 
 func isDML(statement string) bool {
 	return dmlRegex.Match([]byte(statement))
 }
+
+func distinctKind(kindMap map[statementKind]uint64, kind statementKind) bool {
+	target := kindMap[kind]
+
+	var total uint64
+	for k := range kindMap {
+		total = total + kindMap[k]
+	}
+
+	return target == total
+}
+
+func getStatementKind(statement string) statementKind {
+	if isPartitionedDMLOnly(statement) {
+		return statementKindPartitionedDML
+	}
+
+	if isDMLAny(statement) {
+		return statementKindDML
+	}
+
+	return statementKindDDL
+}
+
+func isPartitionedDMLOnly(statement string) bool {
+	return isDMLAny(statement) && !notPartitionedDmlRegex.Match([]byte(statement))
+}
+
+func isDMLAny(statement string) bool {
+	return dmlRegex.Match([]byte(statement))
+}

pkg/spanner/migration_test.go

@@ -17,17 +17,21 @@
 // IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 // CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-package spanner_test
+package spanner
 
 import (
 	"path/filepath"
 	"testing"
+)
 
-	"github.com/cloudspannerecosystem/wrench/pkg/spanner"
+const (
+	TestStmtDDL            = "ALTER TABLE Singers ADD COLUMN Foo STRING(MAX)"
+	TestStmtPartitionedDML = "UPDATE Singers SET FirstName = \"Bar\" WHERE SingerID = \"1\""
+	TestStmtDML            = "INSERT INTO Singers(FirstName) VALUES(\"Bar\")"
 )
 
 func TestLoadMigrations(t *testing.T) {
-	ms, err := spanner.LoadMigrations(filepath.Join("testdata", "migrations"))
+	ms, err := LoadMigrations(filepath.Join("testdata", "migrations"))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -37,29 +41,134 @@ func TestLoadMigrations(t *testing.T) {
 	}
 
 	testcases := []struct {
+		name        string
 		idx         int
 		wantVersion uint
 		wantName    string
 	}{
 		{
+			name:        "want name test and version 2",
 			idx:         0,
 			wantVersion: 2,
 			wantName:    "test",
 		},
 		{
+			name:        "want empty name and version 3",
 			idx:         1,
 			wantVersion: 3,
 			wantName:    "",
 		},
 	}
 
 	for _, tc := range testcases {
-		if ms[tc.idx].Version != tc.wantVersion {
-			t.Errorf("migrations[%d].version want %v, but got %v", tc.idx, tc.wantVersion, ms[tc.idx].Version)
-		}
+		t.Run(tc.name, func(t *testing.T) {
+			if ms[tc.idx].Version != tc.wantVersion {
+				t.Errorf("migrations[%d].version want %v, but got %v", tc.idx, tc.wantVersion, ms[tc.idx].Version)
+			}
+
+			if ms[tc.idx].Name != tc.wantName {
+				t.Errorf("migrations[%d].name want %v, but got %v", tc.idx, tc.wantName, ms[tc.idx].Name)
+			}
+		})
+	}
+}
+
+func Test_getStatementKind(t *testing.T) {
+	tests := []struct {
+		name      string
+		statement string
+		want      statementKind
+	}{
+		{
+			"ALTER statement is DDL",
+			TestStmtDDL,
+			statementKindDDL,
+		},
+		{
+			"UPDATE statement is PartitionedDML",
+			TestStmtPartitionedDML,
+			statementKindPartitionedDML,
+		},
+		{
+			"INSERT statement is DML",
+			TestStmtDML,
+			statementKindDML,
+		},
+		{
+			"lowercase insert statement is DML",
+			TestStmtDML,
+			statementKindDML,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := getStatementKind(tt.statement); got != tt.want {
+				t.Errorf("getStatementKind() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
 
-		if ms[tc.idx].Name != tc.wantName {
-			t.Errorf("migrations[%d].name want %v, but got %v", tc.idx, tc.wantName, ms[tc.idx].Name)
-		}
+func Test_inspectStatementsKind(t *testing.T) {
+	tests := []struct {
+		name       string
+		statements []string
+		want       statementKind
+		wantErr    bool
+	}{
+		{
+			"Only DDL returns DDL",
+			[]string{TestStmtDDL, TestStmtDDL},
+			statementKindDDL,
+			false,
+		},
+		{
+			"Only PartitionedDML returns PartitionedDML",
+			[]string{TestStmtPartitionedDML, TestStmtPartitionedDML},
+			statementKindPartitionedDML,
+			false,
+		},
+		{
+			"Only DML returns DML",
+			[]string{TestStmtDML, TestStmtDML},
+			statementKindDDL,
+			false,
+		},
+		{
+			"DML and DDL returns error",
+			[]string{TestStmtDDL, TestStmtDML},
+			"",
+			true,
+		},
+		{
+			"DML and PartitionedDML returns error",
+			[]string{TestStmtDML, TestStmtPartitionedDML},
+			"",
+			true,
+		},
+		{
+			"DDL and PartitionedDML returns error",
+			[]string{TestStmtDDL, TestStmtPartitionedDML},
+			"",
+			true,
+		},
+		{
+			"no statements defaults to DDL as before",
+			[]string{},
+			statementKindDDL,
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := inspectStatementsKind(tt.statements)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("inspectStatementsKind() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("inspectStatementsKind() got = %v, want %v", got, tt.want)
+			}
+		})
 	}
 }