feat: add sample code snippets for common Spanner and gorm features

samples/emulator/emulator_runner.go

@@ -45,7 +45,7 @@ var containerId string
 // 4. Stop the Docker container with the emulator.
 func RunSampleOnEmulator(sample func(string, string, string) error, ddlStatements ...string) {
 	var err error
-	if err = startEmulator(); err != nil {
+	if _, _, err = startEmulator(); err != nil {
 		log.Fatalf("failed to start emulator: %v", err)
 	}
 	projectId, instanceId, databaseId := "my-project", "my-instance", "my-database"
@@ -64,42 +64,39 @@ func RunSampleOnEmulator(sample func(string, string, string) error, ddlStatement
 	}
 }
 
-func startEmulator() error {
+func startEmulator() (host, port string, err error) {
 	ctx := context.Background()
-	if err := os.Setenv("SPANNER_EMULATOR_HOST", "localhost:9010"); err != nil {
-		return err
-	}
 
 	// Initialize a Docker client.
-	var err error
 	cli, err = client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
 	if err != nil {
-		return err
+		return "", "", err
 	}
 	// Pull the Spanner Emulator docker image.
 	reader, err := cli.ImagePull(ctx, "gcr.io/cloud-spanner-emulator/emulator", image.PullOptions{})
 	if err != nil {
-		return err
+		return "", "", err
 	}
 	defer func() { _ = reader.Close() }()
 	// cli.ImagePull is asynchronous.
 	// The reader needs to be read completely for the pull operation to complete.
 	if _, err := io.Copy(io.Discard, reader); err != nil {
-		return err
+		return "", "", err
 	}
 	// Create and start a container with the emulator.
 	resp, err := cli.ContainerCreate(ctx, &container.Config{
 		Image:        "gcr.io/cloud-spanner-emulator/emulator",
 		ExposedPorts: nat.PortSet{"9010": {}},
 	}, &container.HostConfig{
-		PortBindings: map[nat.Port][]nat.PortBinding{"9010": {{HostIP: "0.0.0.0", HostPort: "9010"}}},
+		AutoRemove:   true,
+		PortBindings: map[nat.Port][]nat.PortBinding{"9010": {{HostIP: "0.0.0.0", HostPort: ""}}},
 	}, nil, nil, "")
 	if err != nil {
-		return err
+		return "", "", err
 	}
 	containerId = resp.ID
 	if err := cli.ContainerStart(ctx, containerId, container.StartOptions{}); err != nil {
-		return err
+		return "", "", err
 	}
 	// Wait max 10 seconds or until the emulator is running.
 	for c := 0; c < 20; c++ {
@@ -108,14 +105,21 @@ func startEmulator() error {
 		<-time.After(500 * time.Millisecond)
 		resp, err := cli.ContainerInspect(ctx, containerId)
 		if err != nil {
-			return fmt.Errorf("failed to inspect container state: %v", err)
+			return "", "", fmt.Errorf("failed to inspect container state: %v", err)
 		}
 		if resp.State.Running {
+			host = resp.NetworkSettings.Ports["9010/tcp"][0].HostIP
+			port = resp.NetworkSettings.Ports["9010/tcp"][0].HostPort
 			break
 		}
 	}
-
-	return nil
+	if host == "" || port == "" {
+		return "", "", fmt.Errorf("emulator did not start successfully")
+	}
+	if err := os.Setenv("SPANNER_EMULATOR_HOST", fmt.Sprintf("%s:%s", host, port)); err != nil {
+		return "", "", err
+	}
+	return
 }
 
 func createInstance(projectId, instanceId string) error {

samples/run_sample.go

@@ -15,14 +15,71 @@
 package main
 
 import (
+	_ "embed"
+	"fmt"
 	"os"
+	"strings"
 
 	"github.com/googleapis/go-gorm-spanner/samples/emulator"
 	samples "github.com/googleapis/go-gorm-spanner/samples/sample_application"
+	"github.com/googleapis/go-gorm-spanner/samples/snippets"
 )
 
+//go:embed snippets/sample_model/data_model.sql
+var createDataModelSQL string
+
 func main() {
-	emulator.RunSampleOnEmulator(func(project string, instance string, database string) error {
-		return samples.RunSample(os.Stdout, "projects/"+project+"/instances/"+instance+"/databases/"+database)
+	// Run the larger sample application.
+	if len(os.Args) == 1 {
+		emulator.RunSampleOnEmulator(func(project string, instance string, database string) error {
+			return samples.RunSample(os.Stdout, "projects/"+project+"/instances/"+instance+"/databases/"+database)
+		})
+		return
+	}
+
+	// Get the DDL statements for the sample data model.
+	ddlStatements := strings.FieldsFunc(createDataModelSQL, func(r rune) bool {
+		return r == ';'
 	})
+	// Skip the last (empty) statement.
+	ddlStatements = ddlStatements[0 : len(ddlStatements)-1]
+
+	// Run one of the sample snippets.
+	sample := os.Args[1]
+
+	switch sample {
+	case "hello_world":
+		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":
+		emulator.RunSampleOnEmulator(snippets.InterleavedTables, ddlStatements...)
+	case "read_only_transaction":
+		emulator.RunSampleOnEmulator(snippets.ReadOnlyTransaction, ddlStatements...)
+	case "read_write_transaction":
+		emulator.RunSampleOnEmulator(snippets.ReadWriteTransaction, ddlStatements...)
+	case "aborted_transaction":
+		emulator.RunSampleOnEmulator(snippets.AbortedTransaction, ddlStatements...)
+	case "migrations":
+		emulator.RunSampleOnEmulator(snippets.Migrations)
+	case "client_library":
+		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/aborted_transaction.go

@@ -0,0 +1,108 @@
+// 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"
+	"gorm.io/gorm"
+)
+
+// AbortedTransaction shows how transaction retries work on Spanner.
+// Read/write transactions guarantee the consistency and atomicity of multiple queries
+// and updates on Spanner. Read/write transactions take locks on the rows that are read
+// and updated. Spanner can abort any read/write transaction due to lock conflicts or
+// due to transient failures (e.g. network errors, machine restarts, etc.).
+//
+// Transactions that fail with an Aborted error should be retried. The Spanner gorm
+// dialect provides the helper function `spannergorm.RunTransaction`
+// for this. It is recommended to run all read/write transactions using this helper
+// function, or add a similar retry function to your own application.
+//
+// Execute the sample with the command `go run run_sample.go aborted_transaction`
+// from the samples directory.
+func AbortedTransaction(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)
+	}
+
+	// Insert a test row.
+	if err := insertVenue(db); err != nil {
+		return err
+	}
+
+	// RunTransaction automatically retries the transaction if it
+	// is aborted by Spanner. It is recommended to use this helper function for
+	// all read/write transactions.
+	attempt := 0
+	if err := spannergorm.RunTransaction(context.Background(), db, func(tx *gorm.DB) error {
+		attempt++
+		fmt.Printf("Executing attempt %d of the first transaction\n", attempt)
+		// Select the venue row in this transaction.
+		var venue sample_model.Venue
+		if err := tx.First(&venue).Error; err != nil {
+			return err
+		}
+		if attempt == 1 {
+			// Execute another read/write transaction that reads and updates the same row.
+			// This will cause this transaction to be aborted by Spanner.
+			if err := readAndUpdateVenueInTransaction(db); err != nil {
+				return err
+			}
+		}
+		venue.Name = venue.Name + " - Updated in first transaction"
+		if err := tx.Updates(&venue).Error; err != nil {
+			return err
+		}
+
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	fmt.Printf("First transaction succeeded after %d attempt(s)\n", attempt)
+
+	return nil
+}
+
+func readAndUpdateVenueInTransaction(db *gorm.DB) error {
+	attempt := 0
+	if err := spannergorm.RunTransaction(context.Background(), db, func(tx *gorm.DB) error {
+		attempt++
+		fmt.Printf("Executing attempt %d of the second transaction\n", attempt)
+		var venue sample_model.Venue
+		if err := tx.First(&venue).Error; err != nil {
+			return err
+		}
+		venue.Name = venue.Name + " - Updated in second transaction"
+		if err := tx.Updates(&venue).Error; err != nil {
+			return err
+		}
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	fmt.Printf("Second transaction succeeded after %d attempt(s)\n", attempt)
+
+	return nil
+}

samples/snippets/auto_save_associations.go

@@ -0,0 +1,102 @@
+// 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 (
+	"database/sql"
+	"fmt"
+
+	"cloud.google.com/go/civil"
+	"cloud.google.com/go/spanner"
+	spannergorm "github.com/googleapis/go-gorm-spanner"
+	"github.com/googleapis/go-gorm-spanner/samples/snippets/sample_model"
+	"gorm.io/gorm"
+)
+
+// AutoSaveAssociations shows how to create a model with one or more associated
+// models in one Create call. gorm uses an insert-or-update statement for these
+// calls. Spanner only supports INSERT OR UPDATE when *ALL* columns are updated.
+// gorm by default only updates the foreign key value when auto-saving associations.
+// You can work around this by calling `db.Session(&gorm.Session{FullSaveAssociations: true})`
+// before saving the model.
+//
+// Execute the sample with the command `go run run_sample.go auto_save_associations`
+// from the samples directory.
+func AutoSaveAssociations(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)
+	}
+
+	// Insert a singer and a few albums.
+	// gorm allows us to create these in one go by creating the model hierarchy
+	// directly in code, and then submitting the top-level model to the Create
+	// function. gorm by default generates a statement that automatically updates
+	// the foreign key value if the child row already exists. This is not supported
+	// on Spanner and results in the following error:
+	// 'spanner only supports UpdateAll or DoNothing for OnConflict clauses'
+	//
+	// This can be worked around by instructing gorm to generate a statement that
+	// updates *ALL* columns of the associated record if it already exists.
+	singer := sample_model.Singer{
+		FirstName: sql.NullString{String: "Angel", Valid: true},
+		LastName:  "Woodward",
+		Active:    true,
+		Albums: []sample_model.Album{
+			{
+				Title:       "Fine Stuff",
+				ReleaseDate: spanner.NullDate{Date: civil.Date{Year: 2024, Month: 11, Day: 11}, Valid: true},
+			},
+			{
+				Title:       "Better Things",
+				ReleaseDate: spanner.NullDate{Date: civil.Date{Year: 2023, Month: 1, Day: 30}, Valid: true},
+			},
+			{
+				Title:       "All Good",
+				ReleaseDate: spanner.NullDate{Date: civil.Date{Year: 2022, Month: 5, Day: 5}, Valid: true},
+			},
+		},
+	}
+	// 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 = db.Session(&gorm.Session{FullSaveAssociations: true}).Create(&singer)
+	if db.Error != nil {
+		return db.Error
+	}
+
+	// Note that gorm only returns the number of affected rows for the top-level
+	// record, i.e. the number of singers that were inserted.
+	fmt.Printf("Inserted %d singer\n", db.RowsAffected)
+
+	// By loading the singer from the database again, we can see that the albums
+	// were also added to the database.
+	db = db.Debug().Preload("Albums").Find(&singer)
+	if db.Error != nil {
+		return db.Error
+	}
+	fmt.Printf("Singer %s has %d albums\n", singer.FullName, len(singer.Albums))
+
+	return nil
+}