From 5d8f2e0ec3cdb0facf923f536827b8cef9172146 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jean-Philippe=20Qu=C3=A9m=C3=A9ner?= Date: Mon, 19 May 2025 13:05:49 +0200 Subject: [PATCH] refactor(unified-storage): extract list iterator and add tests (#105600) --- pkg/storage/unified/sql/backend.go | 67 ----- pkg/storage/unified/sql/list_iterator.go | 73 +++++ pkg/storage/unified/sql/list_iterator_test.go | 256 ++++++++++++++++++ 3 files changed, 329 insertions(+), 67 deletions(-) create mode 100644 pkg/storage/unified/sql/list_iterator.go create mode 100644 pkg/storage/unified/sql/list_iterator_test.go diff --git a/pkg/storage/unified/sql/backend.go b/pkg/storage/unified/sql/backend.go index 15281b9dfcb..71069d5c6d0 100644 --- a/pkg/storage/unified/sql/backend.go +++ b/pkg/storage/unified/sql/backend.go @@ -585,73 +585,6 @@ func (b *backend) ListIterator(ctx context.Context, req *resourcepb.ListRequest, return b.listLatest(ctx, req, cb) } -type listIter struct { - rows db.Rows - offset int64 - listRV int64 - sortAsc bool - - // any error - err error - - // The row - guid string - rv int64 - value []byte - namespace string - resource string - group string - name string - folder string -} - -// ContinueToken implements resource.ListIterator. -func (l *listIter) ContinueToken() string { - return resource.ContinueToken{ResourceVersion: l.listRV, StartOffset: l.offset, SortAscending: l.sortAsc}.String() -} - -func (l *listIter) ContinueTokenWithCurrentRV() string { - return resource.ContinueToken{ResourceVersion: l.rv, StartOffset: l.offset, SortAscending: l.sortAsc}.String() -} - -func (l *listIter) Error() error { - return l.err -} - -func (l *listIter) Name() string { - return l.name -} - -func (l *listIter) Namespace() string { - return l.namespace -} - -func (l *listIter) Folder() string { - return l.folder -} - -// ResourceVersion implements resource.ListIterator. -func (l *listIter) ResourceVersion() int64 { - return l.rv -} - -// Value implements resource.ListIterator. -func (l *listIter) Value() []byte { - return l.value -} - -// Next implements resource.ListIterator. -func (l *listIter) Next() bool { - if l.rows.Next() { - l.offset++ - l.err = l.rows.Scan(&l.guid, &l.rv, &l.namespace, &l.resource, &l.group, &l.name, &l.folder, &l.value) - return true - } - return false -} - -var _ resource.ListIterator = (*listIter)(nil) - // listLatest fetches the resources from the resource table. func (b *backend) listLatest(ctx context.Context, req *resourcepb.ListRequest, cb func(resource.ListIterator) error) (int64, error) { ctx, span := b.tracer.Start(ctx, tracePrefix+"listLatest") diff --git a/pkg/storage/unified/sql/list_iterator.go b/pkg/storage/unified/sql/list_iterator.go new file mode 100644 index 00000000000..96642bf58f3 --- /dev/null +++ b/pkg/storage/unified/sql/list_iterator.go @@ -0,0 +1,73 @@ +package sql + +import ( + "github.com/grafana/grafana/pkg/storage/unified/resource" + "github.com/grafana/grafana/pkg/storage/unified/sql/db" +) + +var _ resource.ListIterator = (*listIter)(nil) + +type listIter struct { + rows db.Rows + offset int64 + listRV int64 + sortAsc bool + + // any error + err error + + // The row + guid string + rv int64 + value []byte + namespace string + resource string + group string + name string + folder string +} + +// ContinueToken implements resource.ListIterator. +func (l *listIter) ContinueToken() string { + return resource.ContinueToken{ResourceVersion: l.listRV, StartOffset: l.offset, SortAscending: l.sortAsc}.String() +} + +func (l *listIter) ContinueTokenWithCurrentRV() string { + return resource.ContinueToken{ResourceVersion: l.rv, StartOffset: l.offset, SortAscending: l.sortAsc}.String() +} + +func (l *listIter) Error() error { + return l.err +} + +func (l *listIter) Name() string { + return l.name +} + +func (l *listIter) Namespace() string { + return l.namespace +} + +func (l *listIter) Folder() string { + return l.folder +} + +// ResourceVersion implements resource.ListIterator. +func (l *listIter) ResourceVersion() int64 { + return l.rv +} + +// Value implements resource.ListIterator. +func (l *listIter) Value() []byte { + return l.value +} + +// Next implements resource.ListIterator. +func (l *listIter) Next() bool { + if l.rows.Next() { + l.offset++ + l.err = l.rows.Scan(&l.guid, &l.rv, &l.namespace, &l.group, &l.resource, &l.name, &l.folder, &l.value) + return true + } + return false +} diff --git a/pkg/storage/unified/sql/list_iterator_test.go b/pkg/storage/unified/sql/list_iterator_test.go new file mode 100644 index 00000000000..ae827237787 --- /dev/null +++ b/pkg/storage/unified/sql/list_iterator_test.go @@ -0,0 +1,256 @@ +package sql + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/grafana/grafana/pkg/infra/db" + "github.com/grafana/grafana/pkg/infra/tracing" + "github.com/grafana/grafana/pkg/setting" + "github.com/grafana/grafana/pkg/storage/unified/resource" + "github.com/grafana/grafana/pkg/storage/unified/resourcepb" + dbsql "github.com/grafana/grafana/pkg/storage/unified/sql/db" + "github.com/grafana/grafana/pkg/storage/unified/sql/db/dbimpl" + "github.com/grafana/grafana/pkg/storage/unified/sql/dbutil" + "github.com/grafana/grafana/pkg/storage/unified/sql/sqltemplate" + "github.com/grafana/grafana/pkg/tests/testsuite" +) + +func TestMain(m *testing.M) { + testsuite.Run(m) +} +func TestListIter(t *testing.T) { + ctx := context.Background() + + grafanaDB := db.InitTestDB(t) + + resourceDBProvider, err := dbimpl.ProvideResourceDB(grafanaDB, setting.NewCfg(), tracing.NewNoopTracerService()) + require.NoError(t, err) + + resourceDB, err := resourceDBProvider.Init(ctx) + require.NoError(t, err) + + dialect := sqltemplate.DialectForDriver(resourceDB.DriverName()) + + testData := []struct { + guid string + resourceVersion int64 + namespace string + resource string + group string + name string + folder string + value []byte + }{ + { + guid: "guid-1", + resourceVersion: 100, + namespace: "namespace-1", + resource: "resource-1", + group: "group-1", + name: "name-1", + folder: "folder-1", + value: []byte(`{"test":"value-1"}`), + }, + { + guid: "guid-2", + resourceVersion: 200, + namespace: "namespace-2", + resource: "resource-2", + group: "group-2", + name: "name-2", + folder: "folder-2", + value: []byte(`{"test":"value-2"}`), + }, + } + + // Insert the test data directly with SQL to include resource_version + err = resourceDB.WithTx(ctx, nil, func(ctx context.Context, tx dbsql.Tx) error { + for _, item := range testData { + _, err := dbutil.Exec(ctx, tx, sqlResourceInsert, sqlResourceRequest{ + SQLTemplate: sqltemplate.New(dialect), + GUID: item.guid, + Folder: item.folder, + ResourceVersion: item.resourceVersion, + WriteEvent: resource.WriteEvent{ + Key: &resourcepb.ResourceKey{ + Namespace: item.namespace, + Resource: item.resource, + Group: item.group, + Name: item.name, + }, + Value: item.value, + PreviousRV: 0, + }, + }) + if err != nil { + return fmt.Errorf("failed to insert test data: %w", err) + } + _, err = dbutil.Exec(ctx, tx, sqlResourceUpdate, sqlResourceRequest{ + SQLTemplate: sqltemplate.New(dialect), + GUID: item.guid, + ResourceVersion: item.resourceVersion, + Folder: item.folder, + WriteEvent: resource.WriteEvent{ + Key: &resourcepb.ResourceKey{ + Namespace: item.namespace, + Resource: item.resource, + Group: item.group, + Name: item.name, + }, + Value: item.value, + PreviousRV: item.resourceVersion, + Type: 1, + }, + }) + if err != nil { + return fmt.Errorf("failed to insert resource version: %w", err) + } + } + return err + }) + require.NoError(t, err) + + t.Run("Next() iterates through results", func(t *testing.T) { + listReq := sqlResourceListRequest{ + SQLTemplate: sqltemplate.New(dialect), + Request: new(resourcepb.ListRequest), + } + rows, err := dbutil.QueryRows(ctx, resourceDB, sqlResourceList, listReq) + require.NoError(t, err) + + iter := &listIter{ + rows: rows, + listRV: 300, + sortAsc: true, + } + + // First row. + require.True(t, iter.Next()) + require.NoError(t, iter.Error()) + require.Equal(t, "guid-1", iter.guid) + require.Equal(t, int64(100), iter.ResourceVersion()) + require.Equal(t, "namespace-1", iter.Namespace()) + require.Equal(t, "resource-1", iter.resource) + require.Equal(t, "group-1", iter.group) + require.Equal(t, "name-1", iter.Name()) + require.Equal(t, "folder-1", iter.Folder()) + require.Equal(t, []byte(`{"test":"value-1"}`), iter.Value()) + + // Second row. + require.True(t, iter.Next()) + require.NoError(t, iter.Error()) + require.Equal(t, "guid-2", iter.guid) + require.Equal(t, int64(200), iter.ResourceVersion()) + require.Equal(t, "namespace-2", iter.Namespace()) + require.Equal(t, "resource-2", iter.resource) + require.Equal(t, "group-2", iter.group) + require.Equal(t, "name-2", iter.Name()) + require.Equal(t, "folder-2", iter.Folder()) + require.Equal(t, []byte(`{"test":"value-2"}`), iter.Value()) + + // No more rows. + require.False(t, iter.Next()) + require.NoError(t, iter.Error()) + }) + + t.Run("Next() returns false when no rows", func(t *testing.T) { + listReq := sqlResourceListRequest{ + SQLTemplate: sqltemplate.New(dialect), + Request: &resourcepb.ListRequest{ + Options: &resourcepb.ListOptions{ + Key: &resourcepb.ResourceKey{ + // Add a filter for a namespace that doesn't exist. + Namespace: "non-existent-namespace", + }, + }, + }, + } + rows, err := dbutil.QueryRows(ctx, resourceDB, sqlResourceList, listReq) + require.NoError(t, err) + + iter := &listIter{ + rows: rows, + listRV: 300, + sortAsc: true, + } + + require.False(t, iter.Next()) + require.NoError(t, iter.Error()) + }) + + t.Run("ContinueToken returns encoded token", func(t *testing.T) { + listReq := sqlResourceListRequest{ + SQLTemplate: sqltemplate.New(dialect), + Request: new(resourcepb.ListRequest), + } + + rows, err := dbutil.QueryRows(ctx, resourceDB, sqlResourceList, listReq) + require.NoError(t, err) + + iter := &listIter{ + rows: rows, + listRV: 300, + sortAsc: true, + } + + require.True(t, iter.Next()) + + token := iter.ContinueToken() + + var actual resource.ContinueToken + b, err := base64.StdEncoding.DecodeString(token) + require.NoError(t, err) + + err = json.Unmarshal(b, &actual) + require.NoError(t, err) + + expected := resource.ContinueToken{ + ResourceVersion: 300, + StartOffset: 1, + SortAscending: true, + } + + require.Equal(t, expected, actual) + }) + + t.Run("ContinueTokenWithCurrentRV uses current row's RV", func(t *testing.T) { + listReq := sqlResourceListRequest{ + SQLTemplate: sqltemplate.New(dialect), + Request: new(resourcepb.ListRequest), + } + + rows, err := dbutil.QueryRows(ctx, resourceDB, sqlResourceList, listReq) + require.NoError(t, err) + + iter := &listIter{ + rows: rows, + listRV: 300, + sortAsc: true, + } + + require.True(t, iter.Next()) + + token := iter.ContinueTokenWithCurrentRV() + + var actual resource.ContinueToken + b, err := base64.StdEncoding.DecodeString(token) + require.NoError(t, err) + + err = json.Unmarshal(b, &actual) + require.NoError(t, err) + + expected := resource.ContinueToken{ + ResourceVersion: 100, + StartOffset: 1, + SortAscending: true, + } + + require.Equal(t, expected, actual) + }) +}