docs: more samples

retry.go

@@ -32,14 +32,12 @@ import (
 // gorm database, and retries the transaction if it is aborted by Spanner.
 func RunTransaction(ctx context.Context, db *gorm.DB, fc func(tx *gorm.DB) error, opts ...*sql.TxOptions) error {
 	// Disable internal (checksum-based) retries on the Spanner database/SQL connection.
-	var opt *sql.TxOptions
 	// Note: gorm also only uses the first option, so it is safe to pick just the first element in the slice.
 	if len(opts) > 0 {
-		opt = opts[0]
+		opts[0].Isolation = spannerdriver.WithDisableRetryAborts(opts[0].Isolation)
 	}
-	opt.Isolation = spannerdriver.WithDisableRetryAborts(opt.Isolation)
 	for {
-		err := db.Transaction(fc, opt)
+		err := db.Transaction(fc, opts...)
 		if err == nil {
 			return nil
 		}

samples/run_sample.go

@@ -52,8 +52,14 @@ func main() {
 		emulator.RunSampleOnEmulator(snippets.HelloWorld, ddlStatements...)
 	case "insert_data":
 		emulator.RunSampleOnEmulator(snippets.InsertData, ddlStatements...)
+	case "upsert":
+		emulator.RunSampleOnEmulator(snippets.Upsert, ddlStatements...)
 	case "batch_insert":
 		emulator.RunSampleOnEmulator(snippets.CreateInBatches, ddlStatements...)
+	case "find_in_batches":
+		emulator.RunSampleOnEmulator(snippets.FindInBatches, ddlStatements...)
+	case "batch_dml":
+		emulator.RunSampleOnEmulator(snippets.BatchDml, ddlStatements...)
 	case "auto_save_associations":
 		emulator.RunSampleOnEmulator(snippets.AutoSaveAssociations, ddlStatements...)
 	case "interleaved_tables":
@@ -70,6 +76,8 @@ func main() {
 		emulator.RunSampleOnEmulator(snippets.ClientLibrary, ddlStatements...)
 	case "uuid_primary_key":
 		emulator.RunSampleOnEmulator(snippets.UuidPrimaryKey)
+	case "bit_reversed_sequence":
+		emulator.RunSampleOnEmulator(snippets.BitReversedSequence)
 	default:
 		fmt.Printf("unknown sample: %s\n", sample)
 		os.Exit(1)

samples/snippets/batch_dml.go

@@ -0,0 +1,77 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package snippets
+
+import (
+	"context"
+	"fmt"
+
+	spannergorm "github.com/googleapis/go-gorm-spanner"
+	"github.com/googleapis/go-gorm-spanner/samples/snippets/sample_model"
+	_ "github.com/googleapis/go-sql-spanner"
+	"gorm.io/gorm"
+)
+
+// BatchDml shows how to use the START BATCH DML / RUN BATCH feature in Spanner to buffer
+// multiple update statements and execute these in one round-trip to Spanner.
+//
+// Execute the sample with the command `go run run_sample.go batch_dml` from this directory.
+func BatchDml(projectId, instanceId, databaseId string) error {
+	db, err := gorm.Open(spannergorm.New(spannergorm.Config{
+		DriverName: "spanner",
+		DSN:        fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectId, instanceId, databaseId),
+	}), &gorm.Config{})
+	if err != nil {
+		return fmt.Errorf("failed to open database connection: %v\n", err)
+	}
+
+	// Insert 50 test singer records.
+	if err := insertTestSingers(db); err != nil {
+		return err
+	}
+
+	// Run a read/write transaction on Spanner that fetches all singers, and updates the Active
+	// flag of each singer as a separate statement.
+	// The START BATCH DML / RUN BATCH statements ensure that these single statements are sent
+	// to Spanner as a single batch. This happens in the following way:
+	// 1. START BATCH DML creates a DML batch on the transaction.
+	// 2. All DML statements on the same transaction are now buffered in memory instead of being sent
+	//    directly to Spanner.
+	// 3. RUN BATCH flushes all statements that have been buffered.
+	var singers []*sample_model.Singer
+	return spannergorm.RunTransaction(context.Background(), db, func(tx *gorm.DB) error {
+		if err := tx.Order("last_name").Find(&singers).Error; err != nil {
+			return err
+		}
+		// Start a DML batch. This will buffer all DML statements in memory until
+		// RUN BATCH is executed.
+		if err := tx.Exec("START BATCH DML").Error; err != nil {
+			return err
+		}
+		for _, singer := range singers {
+			singer.Active = false
+			if err := tx.Save(singer).Error; err != nil {
+				tx.Rollback()
+				return err
+			}
+		}
+		// Execute RUN BATCH. This flushes all buffered DML statements to Spanner.
+		if err := tx.Exec("RUN BATCH").Error; err != nil {
+			return err
+		}
+		fmt.Println("Executed 50 updates in a single DML batch on Spanner")
+		return nil
+	})
+}

samples/snippets/bit_reversed_sequence.go

@@ -0,0 +1,94 @@
+// Copyright 2024 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package snippets
+
+import (
+	"fmt"
+
+	spannergorm "github.com/googleapis/go-gorm-spanner"
+	"gorm.io/gorm"
+)
+
+type blog struct {
+	// gorm automatically assumes that the field with the name ID is the primary
+	// key of the table. It also assumes that it is auto-generated if the type
+	// is an integer of any type.
+	ID    int64
+	Title string
+}
+
+// BitReversedSequence shows how to use a bit-reversed sequence to generate
+// the primary key value in gorm. Using a primary key that is generated by
+// Spanner, means that gorm must use a THEN RETURN clause to return the
+// primary key from the server. As a new primary key value is generated for
+// each transaction attempt, it means that transactions must be defined using
+// a transaction runner, or in some other way must be retried using a retry
+// loop.
+//
+// Execute the sample with the command `go run run_sample.go bit_reversed_sequence`
+// from the samples directory.
+func BitReversedSequence(projectId, instanceId, databaseId string) error {
+	db, err := gorm.Open(spannergorm.New(spannergorm.Config{
+		DriverName: "spanner",
+		DSN:        fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectId, instanceId, databaseId),
+	}), &gorm.Config{PrepareStmt: true})
+	if err != nil {
+		return fmt.Errorf("failed to open database connection: %v\n", err)
+	}
+
+	// Dry-run the migration to verify the DDL statements that will be executed.
+	m := db.Migrator()
+	migrator, ok := m.(spannergorm.SpannerMigrator)
+	if !ok {
+		return fmt.Errorf("unexpected migrator type: %v", m)
+	}
+	// Dry-run the migrations and print the generated statements.
+	statements, err := migrator.AutoMigrateDryRun(&blog{})
+	if err != nil {
+		return fmt.Errorf("could not dry-run migrations: %v", err)
+	}
+	fmt.Print("\nMigrations dry-run generated these statements:\n\n")
+	for _, statement := range statements {
+		fmt.Printf("%s;\n", statement.SQL)
+	}
+	fmt.Println()
+
+	// Create the `blogs` table. This table uses an auto-generated primary key with a backing
+	// bit-reversed sequence.
+	if err := db.Migrator().AutoMigrate(&blog{}); err != nil {
+		return fmt.Errorf("could not execute migrations: %v", err)
+	}
+	// Insert some blog records in the database.
+	if err := db.CreateInBatches([]*blog{
+		{Title: "Blog 1"},
+		{Title: "Blog 2"},
+		{Title: "Blog 3"},
+		{Title: "Blog 4"},
+	}, 100).Error; err != nil {
+		return err
+	}
+
+	// Verify that the records were created.
+	var blogs []blog
+	if err := db.Model(&blog{}).Order("title").Find(&blogs).Error; err != nil {
+		return err
+	}
+	fmt.Println("Found these blogs:")
+	for _, blog := range blogs {
+		fmt.Printf("%d %s\n", blog.ID, blog.Title)
+	}
+
+	return nil
+}

samples/snippets/bind_in_batches.go → samples/snippets/find_in_batches.go

@@ -17,65 +17,66 @@ package snippets
 import (
 	"database/sql"
 	"fmt"
+	"strings"
 
+	"github.com/google/uuid"
 	spannergorm "github.com/googleapis/go-gorm-spanner"
 	"github.com/googleapis/go-gorm-spanner/samples/snippets/sample_model"
 	_ "github.com/googleapis/go-sql-spanner"
 	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
 )
 
 // FindInBatches shows how to process large amounts of data in smaller batches.
 // This reduces the time that locks are held, and ensures that the Spanner
 // transaction mutation limit is not exceeded.
 //
-// Execute the sample with the command `go run run_sample.go hello_world` from this directory.
+// Execute the sample with the command `go run run_sample.go find_in_batches` from this directory.
 func FindInBatches(projectId, instanceId, databaseId string) error {
+	l := logger.Default.LogMode(logger.Info)
 	db, err := gorm.Open(spannergorm.New(spannergorm.Config{
 		DriverName: "spanner",
 		DSN:        fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectId, instanceId, databaseId),
-	}), &gorm.Config{})
+	}), &gorm.Config{Logger: l})
 	if err != nil {
 		return fmt.Errorf("failed to open database connection: %v\n", err)
 	}
 
-	// Create a slice of Singers and insert these in batches of 5.
+	// Insert 50 test singer records.
+	if err := insertTestSingers(db); err != nil {
+		return err
+	}
+
+	// Fetch and process singers in batches of 10. This ensures that locks are not held for longer
+	// than necessary, and that the Spanner transaction mutation limit is not exceeded.
+	// See https://cloud.google.com/spanner/quotas#limits-for for more information on Spanner limits.
+
+	var singers []*sample_model.Singer
+	db.FindInBatches(&singers, 10, func(tx *gorm.DB, batch int) error {
+		for _, singer := range singers {
+			if strings.HasSuffix(singer.FirstName.String, "1") {
+				singer.Active = false
+				if err := tx.Save(singer).Error; err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	})
+
+	return nil
+}
+
+func insertTestSingers(db *gorm.DB) error {
+	// Insert a batch of 50 singers.
 	numSingers := 50
 	singers := make([]sample_model.Singer, 0, numSingers)
 	for i := 0; i < numSingers; i++ {
 		singers = append(singers, sample_model.Singer{
-			FirstName: sql.NullString{String: fmt.Sprintf("First"), Valid: true},
-			LastName:  fmt.Sprintf("Last %d", i),
+			FirstName: sql.NullString{String: fmt.Sprintf("First %d", i), Valid: true},
+			LastName:  uuid.New().String(),
 			Active:    true,
-			Albums: []sample_model.Album{
-				{Title: "Album 1"},
-				{Title: "Album 2"},
-				{Title: "Album 3"},
-				{Title: "Album 4"},
-			},
 		})
 	}
-	// gorm by default tries to only update the association columns when you
-	// auto-create association. This is not supported by Spanner, as Spanner requires
-	// either all columns to be updated, or none (INSERT OR IGNORE).
-	//
-	// By adding `FullSaveAssociations: true` to the session when using auto-save
-	// associations, gorm will generate an INSERT OR UPDATE statement.
-	//
-	// Failing to add `FullSaveAssociations: true` will lead to the following error:
-	// 'spanner only supports UpdateAll or DoNothing for OnConflict clauses'.
-	db.Session(&gorm.Session{FullSaveAssociations: true}).CreateInBatches(&singers, 5)
-	if db.Error != nil {
-		return db.Error
-	}
-
-	// CreateInBatches does not return the affected row count, so we get the number of inserted
-	// singers by executing a query.
-	count := 0
-	db.Raw("select count(1) from singers").Scan(&count)
-	if db.Error != nil {
-		return db.Error
-	}
-	fmt.Printf("Inserted %d singers\n", count)
-
-	return nil
+	return db.CreateInBatches(&singers, 500).Error
 }