feature(spanner): Implement support for DML

Author: jtwatson

database/spanner/README.md

@@ -30,17 +30,11 @@ as described in [README.md#database-urls](../../README.md#database-urls)
 > 1496601752/u add_index_on_user_emails (2m12.155787369s)
 > 1496602638/u create_books_table (2m30.77299181s)
 
-## DDL with comments
+## DDL & DML with comments
 
 At the moment the GCP Spanner backed does not seem to allow for comments (See https://issuetracker.google.com/issues/159730604)
 so in order to be able to use migration with DDL containing comments `x-clean-statements` is required
 
 ## Multiple statements
 
-In order to be able to use more than 1 DDL statement in the same migration file, the file has to be parsed and therefore the `x-clean-statements` flag is required
-
-## Testing
-
-To unit test the `spanner` driver, `SPANNER_DATABASE` needs to be set. You'll
-need to sign-up to Google Cloud Platform (GCP) and have a running Spanner
-instance since it is not possible to run Google Spanner outside GCP.
+In order to be able to use more than 1 DDL or DML statement in the same migration file, the file has to be parsed and therefore the `x-clean-statements` flag is required

database/spanner/examples/migrations2/1481574547_create_users_table.down.sql

@@ -0,0 +1 @@
+DROP TABLE Users;

database/spanner/examples/migrations2/1481574547_create_users_table.up.sql

@@ -0,0 +1,15 @@
+-- Create a table
+CREATE TABLE Users (
+  UserId  INT64,
+  Name    STRING(40),
+  Email   STRING(83)
+) PRIMARY KEY(UserId /* even inline comments */);
+
+CREATE UNIQUE INDEX UsersEmailIndex ON Users (Email);
+
+-- Comments are okay
+
+INSERT INTO Users(UserId, Name, Email)
+  VALUES
+  (100, "Username", "[email protected]"),
+  (200, "Username2", "[email protected]");

database/spanner/examples/migrations2/1496539702_add_city_to_users.down.sql

@@ -0,0 +1 @@
+ALTER TABLE Users DROP COLUMN city;

database/spanner/examples/migrations2/1496539702_add_city_to_users.up.sql

@@ -0,0 +1 @@
+ALTER TABLE Users ADD COLUMN city STRING(100);

database/spanner/spanner.go

@@ -13,8 +13,9 @@ import (
 
 	"cloud.google.com/go/spanner"
 	sdb "cloud.google.com/go/spanner/admin/database/apiv1"
-	"cloud.google.com/go/spanner/spansql"
 
+	"github.com/cloudspannerecosystem/memefish"
+	"github.com/cloudspannerecosystem/memefish/token"
 	"github.com/golang-migrate/migrate/v4"
 	"github.com/golang-migrate/migrate/v4/database"
 
@@ -60,11 +61,9 @@ type Config struct {
 
 // Spanner implements database.Driver for Google Cloud Spanner
 type Spanner struct {
-	db *DB
-
+	db     *DB
 	config *Config
-
-	lock *uatomic.Uint32
+	lock   *uatomic.Uint32
 }
 
 type DB struct {
@@ -179,26 +178,65 @@ func (s *Spanner) Run(migration io.Reader) error {
 		return err
 	}
 
-	stmts := []string{string(migr)}
-	if s.config.CleanStatements {
-		stmts, err = cleanStatements(migr)
-		if err != nil {
-			return err
+	ctx := context.Background()
+
+	if !s.config.CleanStatements {
+		return s.runDdl(ctx, []string{string(migr)})
+	}
+
+	stmtGroups, err := statementGroups(migr)
+	if err != nil {
+		return err
+	}
+
+	for _, group := range stmtGroups {
+		switch group.typ {
+		case statementTypeDDL:
+			if err := s.runDdl(ctx, group.stmts); err != nil {
+				return err
+			}
+		case statementTypeDML:
+			if err := s.runDml(ctx, group.stmts); err != nil {
+				return err
+			}
+		default:
+			return fmt.Errorf("unknown statement type: %s", group.typ)
 		}
 	}
 
-	ctx := context.Background()
+	return nil
+}
+
+func (s *Spanner) runDdl(ctx context.Context, stmts []string) error {
 	op, err := s.db.admin.UpdateDatabaseDdl(ctx, &adminpb.UpdateDatabaseDdlRequest{
 		Database:   s.config.DatabaseName,
 		Statements: stmts,
 	})
 
 	if err != nil {
-		return &database.Error{OrigErr: err, Err: "migration failed", Query: migr}
+		return &database.Error{OrigErr: err, Err: "migration failed", Query: []byte(strings.Join(stmts, ";\n"))}
 	}
 
 	if err := op.Wait(ctx); err != nil {
-		return &database.Error{OrigErr: err, Err: "migration failed", Query: migr}
+		return &database.Error{OrigErr: err, Err: "migration failed", Query: []byte(strings.Join(stmts, ";\n"))}
+	}
+
+	return nil
+}
+
+func (s *Spanner) runDml(ctx context.Context, stmts []string) error {
+	_, err := s.db.data.ReadWriteTransaction(ctx,
+		func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
+			for _, s := range stmts {
+				_, err := txn.Update(ctx, spanner.Statement{SQL: s})
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		})
+	if err != nil {
+		return &database.Error{OrigErr: err, Err: "migration failed", Query: []byte(strings.Join(stmts, ";\n"))}
 	}
 
 	return nil
@@ -345,17 +383,80 @@ func (s *Spanner) ensureVersionTable() (err error) {
 	return nil
 }
 
-func cleanStatements(migration []byte) ([]string, error) {
-	// The Spanner GCP backend does not yet support comments for the UpdateDatabaseDdl RPC
-	// (see https://issuetracker.google.com/issues/159730604) we use
-	// spansql to parse the DDL and output valid stamements without comments
-	ddl, err := spansql.ParseDDL("", string(migration))
-	if err != nil {
-		return nil, err
+type statementType string
+
+const (
+	statementTypeUnknown statementType = ""
+	statementTypeDDL     statementType = "DDL"
+	statementTypeDML     statementType = "DML"
+)
+
+type statementGroup struct {
+	typ   statementType
+	stmts []string
+}
+
+func statementGroups(migr []byte) (groups []*statementGroup, err error) {
+	lex := &memefish.Lexer{
+		File: &token.File{Buffer: string(migr)},
 	}
-	stmts := make([]string, 0, len(ddl.List))
-	for _, stmt := range ddl.List {
-		stmts = append(stmts, stmt.SQL())
+
+	group := &statementGroup{}
+	var stmtTyp statementType
+	var stmt strings.Builder
+	for {
+		if err := lex.NextToken(); err != nil {
+			return nil, err
+		}
+
+		if stmtTyp == statementTypeUnknown {
+			switch {
+			case lex.Token.IsKeywordLike("INSERT") || lex.Token.IsKeywordLike("DELETE") || lex.Token.IsKeywordLike("UPDATE"):
+				stmtTyp = statementTypeDML
+			default:
+				stmtTyp = statementTypeDDL
+			}
+			if group.typ != stmtTyp {
+				if len(group.stmts) > 0 {
+					groups = append(groups, group)
+				}
+				group = &statementGroup{typ: stmtTyp}
+			}
+		}
+
+		if lex.Token.Kind == token.TokenEOF || lex.Token.Kind == ";" {
+			if stmt.Len() > 0 {
+				group.stmts = append(group.stmts, stmt.String())
+			}
+			stmtTyp = statementTypeUnknown
+			stmt.Reset()
+
+			if lex.Token.Kind == token.TokenEOF {
+				if len(group.stmts) > 0 {
+					groups = append(groups, group)
+				}
+
+				break
+			}
+
+			continue
+		}
+
+		if len(lex.Token.Comments) > 0 {
+			// preserve newline where comments are removed
+			if _, err := stmt.WriteString("\n"); err != nil {
+				return nil, err
+			}
+		}
+		if stmt.Len() > 0 {
+			if _, err := stmt.WriteString(lex.Token.Space); err != nil {
+				return nil, err
+			}
+		}
+		if _, err := stmt.WriteString(lex.Token.Raw); err != nil {
+			return nil, err
+		}
 	}
-	return stmts, nil
+
+	return groups, nil
 }