Skip to content

Commit fcd03d1

Browse files
josephschorrjosephschorr
committed
Export the unique, stable, ID from datastore
1 parent e2094f0 commit fcd03d1

23 files changed

+166
-39
lines changed

‎internal/datastore/context.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@ func (p *ctxProxy) MetricsID() (string, error) {
4343
return p.delegate.MetricsID()
4444
}
4545

46+
func (p *ctxProxy) UniqueID(ctx context.Context) (string, error) {
47+
return p.delegate.UniqueID(SeparateContextWithTracing(ctx))
48+
}
49+
4650
func (p *ctxProxy) ReadWriteTx(
4751
ctx context.Context,
4852
f datastore.TxUserFunc,

‎internal/datastore/crdb/crdb.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77
"regexp"
88
"strconv"
9+
"sync/atomic"
910
"time"
1011

1112
"github.com/IBM/pgxpoolprometheus"
@@ -293,6 +294,8 @@ type crdbDatastore struct {
293294
supportsIntegrity bool
294295
expirationEnabled bool
295296
watchEnabled bool
297+
298+
uniqueID atomic.Pointer[string]
296299
}
297300

298301
func (cds *crdbDatastore) SnapshotReader(rev datastore.Revision) datastore.Reader {

‎internal/datastore/crdb/stats.go

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,22 +20,33 @@ const (
2020
colUniqueID = "unique_id"
2121
)
2222

23-
var (
24-
queryReadUniqueID = psql.Select(colUniqueID).From(tableMetadata)
25-
uniqueID string
26-
)
23+
var queryReadUniqueID = psql.Select(colUniqueID).From(tableMetadata)
2724

28-
func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
29-
if len(uniqueID) == 0 {
25+
func (cds *crdbDatastore) UniqueID(ctx context.Context) (string, error) {
26+
if cds.uniqueID.Load() == nil {
3027
sql, args, err := queryReadUniqueID.ToSql()
3128
if err != nil {
32-
return datastore.Stats{}, fmt.Errorf("unable to prepare unique ID sql: %w", err)
29+
return "", fmt.Errorf("unable to prepare unique ID sql: %w", err)
3330
}
31+
32+
var uniqueID string
3433
if err := cds.readPool.QueryRowFunc(ctx, func(ctx context.Context, row pgx.Row) error {
3534
return row.Scan(&uniqueID)
3635
}, sql, args...); err != nil {
37-
return datastore.Stats{}, fmt.Errorf("unable to query unique ID: %w", err)
36+
return "", fmt.Errorf("unable to query unique ID: %w", err)
3837
}
38+
39+
cds.uniqueID.Store(&uniqueID)
40+
return uniqueID, nil
41+
}
42+
43+
return *cds.uniqueID.Load(), nil
44+
}
45+
46+
func (cds *crdbDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
47+
uniqueID, err := cds.UniqueID(ctx)
48+
if err != nil {
49+
return datastore.Stats{}, err
3950
}
4051

4152
var nsDefs []datastore.RevisionedNamespace

‎internal/datastore/memdb/memdb.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,10 @@ func (mdb *memdbDatastore) MetricsID() (string, error) {
108108
return "memdb", nil
109109
}
110110

111+
func (mdb *memdbDatastore) UniqueID(_ context.Context) (string, error) {
112+
return mdb.uniqueID, nil
113+
}
114+
111115
func (mdb *memdbDatastore) SnapshotReader(dr datastore.Revision) datastore.Reader {
112116
mdb.RLock()
113117
defer mdb.RUnlock()

‎internal/datastore/mysql/datastore.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -519,6 +519,8 @@ type Datastore struct {
519519
createTxn sq.InsertBuilder
520520
createBaseTxn string
521521

522+
uniqueID atomic.Pointer[string]
523+
522524
*QueryBuilder
523525
*revisions.CachedOptimizedRevisions
524526
revisions.CommonDecoder
@@ -607,7 +609,7 @@ func (mds *Datastore) isSeeded(ctx context.Context) (bool, error) {
607609
return false, nil
608610
}
609611

610-
_, err = mds.getUniqueID(ctx)
612+
_, err = mds.UniqueID(ctx)
611613
if err != nil {
612614
return false, nil
613615
}

‎internal/datastore/mysql/stats.go

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ func (mds *Datastore) Statistics(ctx context.Context) (datastore.Stats, error) {
3030
}
3131
}
3232

33-
uniqueID, err := mds.getUniqueID(ctx)
33+
uniqueID, err := mds.UniqueID(ctx)
3434
if err != nil {
3535
return datastore.Stats{}, err
3636
}
@@ -88,16 +88,20 @@ func (mds *Datastore) Statistics(ctx context.Context) (datastore.Stats, error) {
8888
}, nil
8989
}
9090

91-
func (mds *Datastore) getUniqueID(ctx context.Context) (string, error) {
92-
sql, args, err := sb.Select(metadataUniqueIDColumn).From(mds.driver.Metadata()).ToSql()
93-
if err != nil {
94-
return "", fmt.Errorf("unable to generate query sql: %w", err)
95-
}
91+
func (mds *Datastore) UniqueID(ctx context.Context) (string, error) {
92+
if mds.uniqueID.Load() == nil {
93+
sql, args, err := sb.Select(metadataUniqueIDColumn).From(mds.driver.Metadata()).ToSql()
94+
if err != nil {
95+
return "", fmt.Errorf("unable to generate query sql: %w", err)
96+
}
9697

97-
var uniqueID string
98-
if err := mds.db.QueryRowContext(ctx, sql, args...).Scan(&uniqueID); err != nil {
99-
return "", fmt.Errorf("unable to query unique ID: %w", err)
98+
var uniqueID string
99+
if err := mds.db.QueryRowContext(ctx, sql, args...).Scan(&uniqueID); err != nil {
100+
return "", fmt.Errorf("unable to query unique ID: %w", err)
101+
}
102+
mds.uniqueID.Store(&uniqueID)
103+
return uniqueID, nil
100104
}
101105

102-
return uniqueID, nil
106+
return *mds.uniqueID.Load(), nil
103107
}

‎internal/datastore/postgres/postgres.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -417,6 +417,7 @@ type pgDatastore struct {
417417
includeQueryParametersInTraces bool
418418

419419
credentialsProvider datastore.CredentialsProvider
420+
uniqueID atomic.Pointer[string]
420421

421422
workerGroup *errgroup.Group
422423
workerCtx context.Context
@@ -698,8 +699,9 @@ func (pgd *pgDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, e
698699
if !state.IsReady {
699700
return state, nil
700701
}
702+
701703
// Ensure a datastore ID is present. This ensures the tables have not been truncated.
702-
uniqueID, err := pgd.datastoreUniqueID(ctx)
704+
uniqueID, err := pgd.UniqueID(ctx)
703705
if err != nil {
704706
return datastore.ReadyState{}, fmt.Errorf("database validation failed: %w; if you have previously run `TRUNCATE`, this database is no longer valid and must be remigrated. See: https://spicedb.dev/d/truncate-unsupported", err)
705707
}

‎internal/datastore/postgres/stats.go

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,16 +29,25 @@ var (
2929
Where(sq.Eq{colRelname: schema.TableTuple})
3030
)
3131

32-
func (pgd *pgDatastore) datastoreUniqueID(ctx context.Context) (string, error) {
33-
idSQL, idArgs, err := queryUniqueID.ToSql()
34-
if err != nil {
35-
return "", fmt.Errorf("unable to generate query sql: %w", err)
32+
func (pgd *pgDatastore) UniqueID(ctx context.Context) (string, error) {
33+
if pgd.uniqueID.Load() == nil {
34+
idSQL, idArgs, err := queryUniqueID.ToSql()
35+
if err != nil {
36+
return "", fmt.Errorf("unable to generate query sql: %w", err)
37+
}
38+
39+
var uniqueID string
40+
if err := pgx.BeginTxFunc(ctx, pgd.readPool, pgd.readTxOptions, func(tx pgx.Tx) error {
41+
return tx.QueryRow(ctx, idSQL, idArgs...).Scan(&uniqueID)
42+
}); err != nil {
43+
return "", fmt.Errorf("unable to query unique ID: %w", err)
44+
}
45+
46+
pgd.uniqueID.Store(&uniqueID)
47+
return uniqueID, nil
3648
}
3749

38-
var uniqueID string
39-
return uniqueID, pgx.BeginTxFunc(ctx, pgd.readPool, pgd.readTxOptions, func(tx pgx.Tx) error {
40-
return tx.QueryRow(ctx, idSQL, idArgs...).Scan(&uniqueID)
41-
})
50+
return *pgd.uniqueID.Load(), nil
4251
}
4352

4453
func (pgd *pgDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {

‎internal/datastore/proxy/checkingreplicated_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,10 @@ func (f fakeDatastore) Statistics(_ context.Context) (datastore.Stats, error) {
141141
return datastore.Stats{}, nil
142142
}
143143

144+
func (f fakeDatastore) UniqueID(_ context.Context) (string, error) {
145+
return "fake", nil
146+
}
147+
144148
func (f fakeDatastore) Close() error {
145149
return nil
146150
}

‎internal/datastore/proxy/indexcheck/fakedatastore.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ func (f fakeDatastore) MetricsID() (string, error) {
2020
return "fake", nil
2121
}
2222

23+
func (f fakeDatastore) UniqueID(_ context.Context) (string, error) {
24+
return "fake", nil
25+
}
26+
2327
func (f fakeDatastore) SnapshotReader(revision datastore.Revision) datastore.Reader {
2428
return fakeSnapshotReader{
2529
revision: revision,

‎internal/datastore/proxy/indexcheck/indexcheck.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ func (p *indexcheckingProxy) MetricsID() (string, error) {
5252
return p.delegate.MetricsID()
5353
}
5454

55+
func (p *indexcheckingProxy) UniqueID(ctx context.Context) (string, error) {
56+
return p.delegate.UniqueID(ctx)
57+
}
58+
5559
func (p *indexcheckingProxy) OptimizedRevision(ctx context.Context) (datastore.Revision, error) {
5660
return p.delegate.OptimizedRevision(ctx)
5761
}

‎internal/datastore/proxy/observable.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,10 @@ func (p *observableProxy) MetricsID() (string, error) {
7272
return p.delegate.MetricsID()
7373
}
7474

75+
func (p *observableProxy) UniqueID(ctx context.Context) (string, error) {
76+
return p.delegate.UniqueID(ctx)
77+
}
78+
7579
func (p *observableProxy) SnapshotReader(rev datastore.Revision) datastore.Reader {
7680
delegateReader := p.delegate.SnapshotReader(rev)
7781
return &observableReader{delegateReader}

‎internal/datastore/proxy/proxy_test/mock.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ type MockDatastore struct {
1717
mock.Mock
1818
}
1919

20+
func (dm *MockDatastore) UniqueID(_ context.Context) (string, error) {
21+
return "mockds", nil
22+
}
23+
2024
func (dm *MockDatastore) SnapshotReader(rev datastore.Revision) datastore.Reader {
2125
args := dm.Called(rev)
2226
return args.Get(0).(datastore.Reader)

‎internal/datastore/proxy/relationshipintegrity.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,10 @@ func (r *relationshipIntegrityProxy) MetricsID() (string, error) {
161161
return r.ds.MetricsID()
162162
}
163163

164+
func (r *relationshipIntegrityProxy) UniqueID(ctx context.Context) (string, error) {
165+
return r.ds.UniqueID(ctx)
166+
}
167+
164168
func (r *relationshipIntegrityProxy) SnapshotReader(rev datastore.Revision) datastore.Reader {
165169
return relationshipIntegrityReader{
166170
parent: r,

‎internal/datastore/proxy/schemacaching/watchingcache_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,10 @@ func (fds *fakeDatastore) MetricsID() (string, error) {
336336
return "fake", nil
337337
}
338338

339+
func (fds *fakeDatastore) UniqueID(_ context.Context) (string, error) {
340+
return "fakedsforwatch", nil
341+
}
342+
339343
func (fds *fakeDatastore) updateNamespace(name string, def *corev1.NamespaceDefinition, revision datastore.Revision) {
340344
fds.lock.Lock()
341345
defer fds.lock.Unlock()

‎internal/datastore/proxy/singleflight.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@ func (p *singleflightProxy) MetricsID() (string, error) {
2828
return p.delegate.MetricsID()
2929
}
3030

31+
func (p *singleflightProxy) UniqueID(ctx context.Context) (string, error) {
32+
return p.delegate.UniqueID(ctx)
33+
}
34+
3135
func (p *singleflightProxy) SnapshotReader(rev datastore.Revision) datastore.Reader {
3236
return p.delegate.SnapshotReader(rev)
3337
}

‎internal/datastore/spanner/revisions.go

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,27 @@ var (
1616
nowStmt = spanner.NewStatement("SELECT CURRENT_TIMESTAMP()")
1717
)
1818

19+
func (sd *spannerDatastore) headRevisionInternal(ctx context.Context) (datastore.Revision, error) {
20+
now, err := sd.now(ctx)
21+
if err != nil {
22+
return datastore.NoRevision, fmt.Errorf(errRevision, err)
23+
}
24+
25+
return revisions.NewForTime(now), nil
26+
}
27+
1928
func (sd *spannerDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) {
29+
return sd.headRevisionInternal(ctx)
30+
}
31+
32+
func (sd *spannerDatastore) now(ctx context.Context) (time.Time, error) {
2033
var timestamp time.Time
2134
if err := sd.client.Single().Query(ctx, nowStmt).Do(func(r *spanner.Row) error {
2235
return r.Columns(&timestamp)
2336
}); err != nil {
24-
return datastore.NoRevision, fmt.Errorf(errRevision, err)
37+
return timestamp, fmt.Errorf(errRevision, err)
2538
}
26-
return revisions.NewForTime(timestamp), nil
39+
return timestamp, nil
2740
}
2841

2942
func (sd *spannerDatastore) staleHeadRevision(ctx context.Context) (datastore.Revision, error) {

‎internal/datastore/spanner/spanner.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"regexp"
99
"strconv"
1010
"sync"
11+
"sync/atomic"
1112
"time"
1213

1314
"cloud.google.com/go/spanner"
@@ -98,6 +99,7 @@ type spannerDatastore struct {
9899

99100
tableSizesStatsTable string
100101
filterMaximumIDCount uint16
102+
uniqueID atomic.Pointer[string]
101103
}
102104

103105
// NewSpannerDatastore returns a datastore backed by cloud spanner

‎internal/datastore/spanner/stats.go

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,16 +28,29 @@ var querySomeRandomRelationships = fmt.Sprintf(`SELECT %s FROM %s LIMIT 10`,
2828

2929
const defaultEstimatedBytesPerRelationships = 20 // determined by looking at some sample clusters
3030

31+
func (sd *spannerDatastore) UniqueID(ctx context.Context) (string, error) {
32+
if sd.uniqueID.Load() == nil {
33+
var uniqueID string
34+
if err := sd.client.Single().Read(
35+
ctx,
36+
tableMetadata,
37+
spanner.AllKeys(),
38+
[]string{colUniqueID},
39+
).Do(func(r *spanner.Row) error {
40+
return r.Columns(&uniqueID)
41+
}); err != nil {
42+
return "", fmt.Errorf("unable to read unique ID: %w", err)
43+
}
44+
sd.uniqueID.Store(&uniqueID)
45+
return uniqueID, nil
46+
}
47+
48+
return *sd.uniqueID.Load(), nil
49+
}
50+
3151
func (sd *spannerDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
32-
var uniqueID string
33-
if err := sd.client.Single().Read(
34-
context.Background(),
35-
tableMetadata,
36-
spanner.AllKeys(),
37-
[]string{colUniqueID},
38-
).Do(func(r *spanner.Row) error {
39-
return r.Columns(&uniqueID)
40-
}); err != nil {
52+
uniqueID, err := sd.UniqueID(ctx)
53+
if err != nil {
4154
return datastore.Stats{}, fmt.Errorf("unable to read unique ID: %w", err)
4255
}
4356

‎pkg/datastore/datastore.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,11 @@ type ReadOnlyDatastore interface {
683683
// and may not be unique; callers should not rely on uniqueness.
684684
MetricsID() (string, error)
685685

686+
// UniqueID returns a unique identifier for the datastore. This identifier
687+
// must be stable across restarts of the datastore if the datastore is
688+
// persistent.
689+
UniqueID(context.Context) (string, error)
690+
686691
// SnapshotReader creates a read-only handle that reads the datastore at the specified revision.
687692
// Any errors establishing the reader will be returned by subsequent calls.
688693
SnapshotReader(Revision) Reader

‎pkg/datastore/datastore_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -604,6 +604,10 @@ type fakeDatastore struct {
604604
delegate Datastore
605605
}
606606

607+
func (f fakeDatastore) UniqueID(_ context.Context) (string, error) {
608+
return "fake", nil
609+
}
610+
607611
func (f fakeDatastore) Unwrap() Datastore {
608612
return f.delegate
609613
}

0 commit comments

Comments
 (0)