Add DML migration processing to spanner

database/spanner/README.md

@@ -11,14 +11,15 @@ The DSN must be given in the following format.
 
 as described in [README.md#database-urls](../../README.md#database-urls)
 
-| Param | WithInstance Config | Description |
-| ----- | ------------------- | ----------- |
-| `x-migrations-table` | `MigrationsTable` | Name of the migrations table |
-| `x-clean-statements` | `CleanStatements` | Whether to parse and clean DDL statements before running migration towards Spanner (Required for comments and multiple statements) |
-| `url` | `DatabaseName` | The full path to the Spanner database resource. If provided as part of `Config` it must not contain a scheme or query string to match the format `projects/{projectId}/instances/{instanceId}/databases/{databaseName}`|
-| `projectId` || The Google Cloud Platform project id
-| `instanceId` || The id of the instance running Spanner
-| `databaseName` || The name of the Spanner database
+| Param                | WithInstance Config                    | Description                                                                                                                                                                                                             |
+|----------------------|----------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `x-migrations-table` | `MigrationsTable`                      | Name of the migrations table                                                                                                                                                                                            |
+| `x-clean-statements` | `CleanStatements`                      | Whether to parse and clean DDL statements before running migration towards Spanner (Required for comments and multiple statements)                                                                                      |
+| `x-dml-comment-flag` | `DmlFlag`                              | Comment flag to treat a migration file as DML                                                                                                                                                                           |
+| `url`                | `DatabaseName`                         | The full path to the Spanner database resource. If provided as part of `Config` it must not contain a scheme or query string to match the format `projects/{projectId}/instances/{instanceId}/databases/{databaseName}` |
+| `projectId`          || The Google Cloud Platform project id   
+| `instanceId`         || The id of the instance running Spanner 
+| `databaseName`       || The name of the Spanner database       
 
 > **Note:** Google Cloud Spanner migrations can take a considerable amount of 
 > time. The migrations provided as part of the example take about 6 minutes to 
@@ -39,6 +40,11 @@ so in order to be able to use migration with DDL containing comments `x-clean-st
 
 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
 
+## DML
+
+In order to have a migration file with DML you need to start the migration file with a one line comment for example #DML 
+and use `x-dml-comment-flag=DML`
+
 ## Testing
 
 To unit test the `spanner` driver, `SPANNER_DATABASE` needs to be set. You'll

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

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

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

@@ -0,0 +1,10 @@
+CREATE TABLE Users (
+  UserId  INT64,
+  Name    STRING(40),
+  Email   STRING(83)
+) PRIMARY KEY(UserId);
+CREATE TABLE Users2 (
+                       UserId  INT64,
+                       Name    STRING(40),
+                       Email   STRING(83)
+) PRIMARY KEY(UserId);

database/spanner/examples/migrationswithdml/1481574548_update_users.down.sql

@@ -0,0 +1,5 @@
+#DML
+UPDATE Users SET Name = null WHERE UserId is null;
+UPDATE Users SET Name = null WHERE UserId is null;
+UPDATE Users2 SET Name = null WHERE UserId is null;
+UPDATE Users2 SET Name = null WHERE UserId is null;

database/spanner/examples/migrationswithdml/1481574548_update_users.up.sql

@@ -0,0 +1,7 @@
+#DML
+UPDATE Users SET Name = 'name'
+             WHERE UserId is null;
+UPDATE Users SET Name = 'name' WHERE UserId is null;
+UPDATE Users2 SET Name = 'name' -- Set the name
+              WHERE UserId is null;
+UPDATE Users2 SET Name = 'name' WHERE UserId is null;

database/spanner/spanner.go

@@ -10,6 +10,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"unicode"
 
 	"cloud.google.com/go/spanner"
 	sdb "cloud.google.com/go/spanner/admin/database/apiv1"
@@ -56,6 +57,9 @@ type Config struct {
 	// Parsing outputs clean DDL statements such as reformatted
 	// and void of comments.
 	CleanStatements bool
+	// Comment flag name to treat a migration file as DML
+	// Example #DML
+	DmlFlag string
 }
 
 // Spanner implements database.Driver for Google Cloud Spanner
@@ -125,9 +129,10 @@ func (s *Spanner) Open(url string) (database.Driver, error) {
 		log.Fatal(err)
 	}
 
+	dmlFlag := purl.Query().Get("x-dml-comment-flag")
 	migrationsTable := purl.Query().Get("x-migrations-table")
-
 	cleanQuery := purl.Query().Get("x-clean-statements")
+
 	clean := false
 	if cleanQuery != "" {
 		clean, err = strconv.ParseBool(cleanQuery)
@@ -141,6 +146,7 @@ func (s *Spanner) Open(url string) (database.Driver, error) {
 		DatabaseName:    dbname,
 		MigrationsTable: migrationsTable,
 		CleanStatements: clean,
+		DmlFlag:         dmlFlag,
 	})
 }
 
@@ -174,26 +180,50 @@ func (s *Spanner) Run(migration io.Reader) error {
 		return err
 	}
 
-	stmts := []string{string(migr)}
-	if s.config.CleanStatements {
-		stmts, err = cleanStatements(migr)
+	migrstr := string(migr)
+	stmts := []string{migrstr}
+
+	ctx := context.Background()
+
+	if s.config.DmlFlag != "" && strings.HasPrefix(migrstr, "#"+s.config.DmlFlag) {
+		_, err := s.db.data.ReadWriteTransaction(ctx,
+			func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
+				for _, v := range strings.Split(migrstr, ";") {
+					if replaceWhiteSpaceWithSpace(strings.TrimSpace(v)) != "" {
+						_, err = txn.Update(ctx, spanner.NewStatement(v))
+						if err != nil {
+							return err
+						}
+					}
+				}
+				return nil
+			})
+
 		if err != nil {
-			return err
+			return &database.Error{OrigErr: err, Err: "migration failed", Query: migr}
 		}
-	}
 
-	ctx := context.Background()
-	op, err := s.db.admin.UpdateDatabaseDdl(ctx, &adminpb.UpdateDatabaseDdlRequest{
-		Database:   s.config.DatabaseName,
-		Statements: stmts,
-	})
+	} else {
 
-	if err != nil {
-		return &database.Error{OrigErr: err, Err: "migration failed", Query: migr}
-	}
+		if s.config.CleanStatements {
+			stmts, err = cleanStatements(migr)
+			if err != nil {
+				return err
+			}
+		}
 
-	if err := op.Wait(ctx); err != nil {
-		return &database.Error{OrigErr: err, Err: "migration failed", Query: migr}
+		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}
+		}
+
+		if err := op.Wait(ctx); err != nil {
+			return &database.Error{OrigErr: err, Err: "migration failed", Query: migr}
+		}
 	}
 
 	return nil
@@ -354,3 +384,13 @@ func cleanStatements(migration []byte) ([]string, error) {
 	}
 	return stmts, nil
 }
+
+func replaceWhiteSpaceWithSpace(str string) string {
+	s := strings.Map(func(r rune) rune {
+		if unicode.IsSpace(r) {
+			return ' '
+		}
+		return r
+	}, str)
+	return strings.TrimSpace(strings.Join(strings.Fields(s), " "))
+}

database/spanner/spanner_test.go

@@ -61,6 +61,22 @@ func TestMigrate(t *testing.T) {
 	})
 }
 
+func TestMigrateWithMultipleDDLStatementsAndDML(t *testing.T) {
+	withSpannerEmulator(t, func(t *testing.T) {
+		s := &Spanner{}
+		uri := fmt.Sprintf("spanner://%s?x-dml-comment-flag=DML&x-clean-statements=true", db)
+		d, err := s.Open(uri)
+		if err != nil {
+			t.Fatal(err)
+		}
+		m, err := migrate.NewWithDatabaseInstance("file://./examples/migrationswithdml", uri, d)
+		if err != nil {
+			t.Fatal(err)
+		}
+		dt.TestMigrate(t, m)
+	})
+}
+
 func TestCleanStatements(t *testing.T) {
 	testCases := []struct {
 		name           string