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

     1  //go:build ci && docker
     2  // +build ci,docker
     3  
     4  package spanner
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"cloud.google.com/go/spanner"
    12  	admin "cloud.google.com/go/spanner/admin/database/apiv1"
    13  	"cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
    14  	"github.com/stretchr/testify/require"
    15  	"google.golang.org/grpc/codes"
    16  	"google.golang.org/grpc/status"
    17  
    18  	testdatastore "github.com/authzed/spicedb/internal/testserver/datastore"
    19  	"github.com/authzed/spicedb/pkg/datastore"
    20  	"github.com/authzed/spicedb/pkg/datastore/test"
    21  	corev1 "github.com/authzed/spicedb/pkg/proto/core/v1"
    22  	"github.com/authzed/spicedb/pkg/tuple"
    23  )
    24  
    25  // Implement TestableDatastore interface
    26  func (sd *spannerDatastore) ExampleRetryableError() error {
    27  	return status.New(codes.Aborted, "retryable").Err()
    28  }
    29  
    30  func TestSpannerDatastore(t *testing.T) {
    31  	ctx := context.Background()
    32  	b := testdatastore.RunSpannerForTesting(t, "", "head")
    33  
    34  	// TODO(jschorr): Once https://github.com/GoogleCloudPlatform/cloud-spanner-emulator/issues/74 has been resolved,
    35  	// change back to `All` to re-enable watch and GC tests.
    36  	// GC tests are disabled because they depend also on the ability to configure change streams with custom retention.
    37  	test.AllWithExceptions(t, test.DatastoreTesterFunc(func(revisionQuantization, _, _ time.Duration, watchBufferLength uint16) (datastore.Datastore, error) {
    38  		ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
    39  			ds, err := NewSpannerDatastore(ctx, uri,
    40  				RevisionQuantization(revisionQuantization),
    41  				WatchBufferLength(watchBufferLength))
    42  			require.NoError(t, err)
    43  			return ds
    44  		})
    45  		return ds, nil
    46  	}), test.WithCategories(test.GCCategory, test.WatchCategory, test.StatsCategory))
    47  
    48  	t.Run("TestFakeStats", createDatastoreTest(
    49  		b,
    50  		FakeStatsTest,
    51  	))
    52  }
    53  
    54  type datastoreTestFunc func(t *testing.T, ds datastore.Datastore)
    55  
    56  func createDatastoreTest(b testdatastore.RunningEngineForTest, tf datastoreTestFunc, options ...Option) func(*testing.T) {
    57  	return func(t *testing.T) {
    58  		ctx := context.Background()
    59  		ds := b.NewDatastore(t, func(engine, uri string) datastore.Datastore {
    60  			ds, err := NewSpannerDatastore(ctx, uri, options...)
    61  			require.NoError(t, err)
    62  			return ds
    63  		})
    64  		defer ds.Close()
    65  
    66  		tf(t, ds)
    67  	}
    68  }
    69  
    70  const createFakeStatsTable = `
    71  CREATE TABLE fake_stats_table (
    72    interval_end TIMESTAMP,
    73    table_name STRING(MAX),
    74    used_bytes INT64,
    75  ) PRIMARY KEY (table_name, interval_end)
    76  `
    77  
    78  func FakeStatsTest(t *testing.T, ds datastore.Datastore) {
    79  	spannerDS := ds.(*spannerDatastore)
    80  	spannerDS.tableSizesStatsTable = "fake_stats_table"
    81  
    82  	spannerClient := spannerDS.client
    83  	ctx := context.Background()
    84  
    85  	adminClient, err := admin.NewDatabaseAdminClient(ctx)
    86  	require.NoError(t, err)
    87  
    88  	// Manually add the stats table to simulate the table that the emulator doesn't create.
    89  	updateOp, err := adminClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{
    90  		Database: spannerClient.DatabaseName(),
    91  		Statements: []string{
    92  			createFakeStatsTable,
    93  		},
    94  	})
    95  	require.NoError(t, err)
    96  
    97  	err = updateOp.Wait(ctx)
    98  	require.NoError(t, err)
    99  
   100  	// Call stats with no stats rows and no relationship rows.
   101  	stats, err := ds.Statistics(ctx)
   102  	require.NoError(t, err)
   103  	require.Equal(t, uint64(0), stats.EstimatedRelationshipCount)
   104  
   105  	// Add some relationships.
   106  	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, tx datastore.ReadWriteTransaction) error {
   107  		return tx.WriteRelationships(ctx, []*corev1.RelationTupleUpdate{
   108  			tuple.Create(tuple.MustParse("document:foo#viewer@user:tom")),
   109  			tuple.Create(tuple.MustParse("document:foo#viewer@user:sarah")),
   110  			tuple.Create(tuple.MustParse("document:foo#viewer@user:fred")),
   111  		})
   112  	})
   113  	require.NoError(t, err)
   114  
   115  	// Call stats with no stats rows and some relationship rows.
   116  	stats, err = ds.Statistics(ctx)
   117  	require.NoError(t, err)
   118  	require.Equal(t, uint64(0), stats.EstimatedRelationshipCount)
   119  
   120  	// Add some stats row with a byte count.
   121  	_, err = spannerClient.Apply(ctx, []*spanner.Mutation{
   122  		spanner.Insert("fake_stats_table", []string{"interval_end", "table_name", "used_bytes"}, []interface{}{
   123  			time.Now().UTC().Add(-100 * time.Second), tableRelationship, 100,
   124  		}),
   125  	})
   126  	require.NoError(t, err)
   127  
   128  	// Call stats with a stats row and some relationship rows and ensure we get an estimate.
   129  	stats, err = ds.Statistics(ctx)
   130  	require.NoError(t, err)
   131  	require.Equal(t, uint64(3), stats.EstimatedRelationshipCount)
   132  
   133  	// Add some more relationships.
   134  	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, tx datastore.ReadWriteTransaction) error {
   135  		return tx.WriteRelationships(ctx, []*corev1.RelationTupleUpdate{
   136  			tuple.Create(tuple.MustParse("document:foo#viewer@user:tommy1236512365123651236512365123612365123655")),
   137  			tuple.Create(tuple.MustParse("document:foo#viewer@user:sara1236512365123651236512365123651236512365")),
   138  			tuple.Create(tuple.MustParse("document:foo#viewer@user:freddy1236512365123651236512365123651236512365")),
   139  		})
   140  	})
   141  	require.NoError(t, err)
   142  
   143  	// Call stats again and ensure it uses the cached relationship size value, even if we'd addded more relationships.
   144  	stats, err = ds.Statistics(ctx)
   145  	require.NoError(t, err)
   146  	require.Equal(t, uint64(3), stats.EstimatedRelationshipCount)
   147  }