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

     1  package migrations
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  
     7  	"cloud.google.com/go/spanner"
     8  	admin "cloud.google.com/go/spanner/admin/database/apiv1"
     9  	"google.golang.org/api/option"
    10  	"google.golang.org/grpc/codes"
    11  
    12  	log "github.com/authzed/spicedb/internal/logging"
    13  	"github.com/authzed/spicedb/pkg/migrate"
    14  )
    15  
    16  const (
    17  	tableSchemaVersion = "schema_version"
    18  	colVersionNum      = "version_num"
    19  	emulatorSettingKey = "SPANNER_EMULATOR_HOST"
    20  )
    21  
    22  // SpannerMigrationDriver can migrate a Cloud Spanner instance
    23  // The adminClient is required for DDL changes
    24  type SpannerMigrationDriver struct {
    25  	client      *spanner.Client
    26  	adminClient *admin.DatabaseAdminClient
    27  }
    28  
    29  // Wrapper makes it possible to forward the spanner clients to the MigrationFunc's to execute
    30  type Wrapper struct {
    31  	client      *spanner.Client
    32  	adminClient *admin.DatabaseAdminClient
    33  }
    34  
    35  // NewSpannerDriver returns a migration driver for the given Cloud Spanner instance
    36  func NewSpannerDriver(ctx context.Context, database, credentialsFilePath, emulatorHost string) (*SpannerMigrationDriver, error) {
    37  	if len(emulatorHost) > 0 {
    38  		err := os.Setenv(emulatorSettingKey, emulatorHost)
    39  		if err != nil {
    40  			return nil, err
    41  		}
    42  	}
    43  	log.Ctx(ctx).Info().Str("spanner-emulator-host", os.Getenv(emulatorSettingKey)).Msg("spanner emulator")
    44  	log.Ctx(ctx).Info().Str("credentials", credentialsFilePath).Str("db", database).Msg("connecting")
    45  	client, err := spanner.NewClient(ctx, database, option.WithCredentialsFile(credentialsFilePath))
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	adminClient, err := admin.NewDatabaseAdminClient(ctx, option.WithCredentialsFile(credentialsFilePath))
    51  	if err != nil {
    52  		return nil, err
    53  	}
    54  
    55  	return &SpannerMigrationDriver{client, adminClient}, nil
    56  }
    57  
    58  // VersionProvider returns the migration version a specific spanner datastore is running at
    59  type VersionProvider interface {
    60  	Version(ctx context.Context) (string, error)
    61  }
    62  
    63  // NewSpannerVersionChecker returns a VersionProvider for the argument spanner.Client
    64  func NewSpannerVersionChecker(c *spanner.Client) VersionProvider {
    65  	return &SpannerMigrationDriver{c, nil}
    66  }
    67  
    68  func (smd *SpannerMigrationDriver) Version(ctx context.Context) (string, error) {
    69  	var schemaRevision string
    70  	iter := smd.client.Single().Read(
    71  		ctx,
    72  		tableSchemaVersion,
    73  		spanner.AllKeys(),
    74  		[]string{colVersionNum},
    75  	)
    76  	defer iter.Stop()
    77  
    78  	if err := iter.Do(func(r *spanner.Row) error {
    79  		return r.Columns(&schemaRevision)
    80  	}); err != nil {
    81  		if spanner.ErrCode(err) == codes.NotFound {
    82  			// There is no schema table, empty database
    83  			return "", nil
    84  		}
    85  		return "", err
    86  	}
    87  
    88  	return schemaRevision, nil
    89  }
    90  
    91  // Conn returns the underlying spanner clients in a Wrapper instance for MigrationFunc to use
    92  func (smd *SpannerMigrationDriver) Conn() Wrapper {
    93  	return Wrapper{client: smd.client, adminClient: smd.adminClient}
    94  }
    95  
    96  func (smd *SpannerMigrationDriver) RunTx(ctx context.Context, f migrate.TxMigrationFunc[*spanner.ReadWriteTransaction]) error {
    97  	_, err := smd.client.ReadWriteTransaction(ctx, func(ctx context.Context, rwt *spanner.ReadWriteTransaction) error {
    98  		return f(ctx, rwt)
    99  	})
   100  	return err
   101  }
   102  
   103  func (smd *SpannerMigrationDriver) WriteVersion(_ context.Context, rwt *spanner.ReadWriteTransaction, version, replaced string) error {
   104  	return rwt.BufferWrite([]*spanner.Mutation{
   105  		spanner.Delete(tableSchemaVersion, spanner.KeySetFromKeys(spanner.Key{replaced})),
   106  		spanner.Insert(tableSchemaVersion, []string{colVersionNum}, []any{version}),
   107  	})
   108  }
   109  
   110  func (smd *SpannerMigrationDriver) Close(_ context.Context) error {
   111  	smd.client.Close()
   112  	return nil
   113  }
   114  
   115  var _ migrate.Driver[Wrapper, *spanner.ReadWriteTransaction] = &SpannerMigrationDriver{}