github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/testserver/datastore/spanner.go (about)

     1  //go:build docker
     2  // +build docker
     3  
     4  package datastore
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"os"
    10  	"testing"
    11  	"time"
    12  
    13  	database "cloud.google.com/go/spanner/admin/database/apiv1"
    14  	adminpb "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
    15  	instances "cloud.google.com/go/spanner/admin/instance/apiv1"
    16  	"cloud.google.com/go/spanner/admin/instance/apiv1/instancepb"
    17  	"github.com/google/uuid"
    18  	"github.com/ory/dockertest/v3"
    19  	"github.com/stretchr/testify/require"
    20  
    21  	"github.com/authzed/spicedb/internal/datastore/spanner/migrations"
    22  	"github.com/authzed/spicedb/pkg/datastore"
    23  	"github.com/authzed/spicedb/pkg/migrate"
    24  	"github.com/authzed/spicedb/pkg/secrets"
    25  )
    26  
    27  type spannerTest struct {
    28  	hostname        string
    29  	targetMigration string
    30  }
    31  
    32  // RunSpannerForTesting returns a RunningEngineForTest for spanner
    33  func RunSpannerForTesting(t testing.TB, bridgeNetworkName string, targetMigration string) RunningEngineForTest {
    34  	pool, err := dockertest.NewPool("")
    35  	require.NoError(t, err)
    36  
    37  	name := fmt.Sprintf("spanner-%s", uuid.New().String())
    38  	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
    39  		Name:         name,
    40  		Repository:   "gcr.io/cloud-spanner-emulator/emulator",
    41  		Tag:          "1.5.11",
    42  		ExposedPorts: []string{"9010/tcp"},
    43  		NetworkID:    bridgeNetworkName,
    44  	})
    45  	require.NoError(t, err)
    46  
    47  	t.Cleanup(func() {
    48  		require.NoError(t, pool.Purge(resource))
    49  	})
    50  
    51  	port := resource.GetPort("9010/tcp")
    52  	spannerEmulatorAddr := fmt.Sprintf("localhost:%s", port)
    53  	require.NoError(t, os.Setenv("SPANNER_EMULATOR_HOST", spannerEmulatorAddr))
    54  
    55  	require.NoError(t, pool.Retry(func() error {
    56  		ctx, cancel := context.WithTimeout(context.Background(), dockerBootTimeout)
    57  		defer cancel()
    58  
    59  		instancesClient, err := instances.NewInstanceAdminClient(ctx)
    60  		if err != nil {
    61  			return err
    62  		}
    63  		defer func() { require.NoError(t, instancesClient.Close()) }()
    64  
    65  		ctx, cancel = context.WithTimeout(context.Background(), dockerBootTimeout)
    66  		defer cancel()
    67  		_, err = instancesClient.CreateInstance(ctx, &instancepb.CreateInstanceRequest{
    68  			Parent:     "projects/fake-project-id",
    69  			InstanceId: "init",
    70  			Instance: &instancepb.Instance{
    71  				Config:      "emulator-config",
    72  				DisplayName: "Test Instance",
    73  				NodeCount:   1,
    74  			},
    75  		})
    76  		return err
    77  	}))
    78  
    79  	builder := &spannerTest{
    80  		targetMigration: targetMigration,
    81  	}
    82  	if bridgeNetworkName != "" {
    83  		builder.hostname = name
    84  	}
    85  
    86  	return builder
    87  }
    88  
    89  func (b *spannerTest) ExternalEnvVars() []string {
    90  	return []string{fmt.Sprintf("SPANNER_EMULATOR_HOST=%s:9010", b.hostname)}
    91  }
    92  
    93  func (b *spannerTest) NewDatabase(t testing.TB) string {
    94  	t.Logf("using spanner emulator, host: %s", os.Getenv("SPANNER_EMULATOR_HOST"))
    95  
    96  	uniquePortion, err := secrets.TokenHex(4)
    97  	require.NoError(t, err)
    98  
    99  	newInstanceName := fmt.Sprintf("fake-instance-%s", uniquePortion)
   100  
   101  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   102  	defer cancel()
   103  
   104  	instancesClient, err := instances.NewInstanceAdminClient(ctx)
   105  	require.NoError(t, err)
   106  	defer instancesClient.Close()
   107  
   108  	createInstanceOp, err := instancesClient.CreateInstance(ctx, &instancepb.CreateInstanceRequest{
   109  		Parent:     "projects/fake-project-id",
   110  		InstanceId: newInstanceName,
   111  		Instance: &instancepb.Instance{
   112  			Config:      "emulator-config",
   113  			DisplayName: "Test Instance",
   114  			NodeCount:   1,
   115  		},
   116  	})
   117  	require.NoError(t, err)
   118  
   119  	spannerInstance, err := createInstanceOp.Wait(ctx)
   120  	require.NoError(t, err)
   121  
   122  	adminClient, err := database.NewDatabaseAdminClient(ctx)
   123  	require.NoError(t, err)
   124  	defer adminClient.Close()
   125  
   126  	dbID := "fake-database-id"
   127  	op, err := adminClient.CreateDatabase(ctx, &adminpb.CreateDatabaseRequest{
   128  		Parent:          spannerInstance.Name,
   129  		CreateStatement: "CREATE DATABASE `" + dbID + "`",
   130  	})
   131  	require.NoError(t, err)
   132  
   133  	db, err := op.Wait(ctx)
   134  	require.NoError(t, err)
   135  	return db.Name
   136  }
   137  
   138  func (b *spannerTest) NewDatastore(t testing.TB, initFunc InitFunc) datastore.Datastore {
   139  	db := b.NewDatabase(t)
   140  
   141  	migrationDriver, err := migrations.NewSpannerDriver(context.Background(), db, "", os.Getenv("SPANNER_EMULATOR_HOST"))
   142  	require.NoError(t, err)
   143  
   144  	err = migrations.SpannerMigrations.Run(context.Background(), migrationDriver, b.targetMigration, migrate.LiveRun)
   145  	require.NoError(t, err)
   146  
   147  	return initFunc("spanner", db)
   148  }