From c8f516fc551abc389f1e7e410d05bf419891ea0c Mon Sep 17 00:00:00 2001 From: Georges Chaudy Date: Mon, 8 Jul 2024 21:55:25 +0200 Subject: [PATCH] pagination first page --- pkg/storage/unified/sql/backend.go | 91 ++++++++++++------- pkg/storage/unified/sql/backend_test.go | 50 ++++++++++ pkg/storage/unified/sql/continue.go | 32 +++++++ .../unified/sql/data/resource_list.sql | 4 + pkg/storage/unified/sql/queries_test.go | 1 + .../testdata/resource_list_mysql_sqlite.sql | 5 +- 6 files changed, 150 insertions(+), 33 deletions(-) create mode 100644 pkg/storage/unified/sql/continue.go diff --git a/pkg/storage/unified/sql/backend.go b/pkg/storage/unified/sql/backend.go index a8a4750e86b1..a12fbbc8b5af 100644 --- a/pkg/storage/unified/sql/backend.go +++ b/pkg/storage/unified/sql/backend.go @@ -138,8 +138,6 @@ func (b *backend) WriteEvent(ctx context.Context, event resource.WriteEvent) (in return b.update(ctx, event) case resource.WatchEvent_DELETED: return b.delete(ctx, event) - default: - } return 0, fmt.Errorf("unsupported event type") } @@ -353,37 +351,64 @@ func (b *backend) PrepareList(ctx context.Context, req *resource.ListRequest) (* _, span := b.tracer.Start(ctx, trace_prefix+"List") defer span.End() - readReq := sqlResourceListRequest{ - SQLTemplate: sqltemplate.New(b.sqlDialect), - Request: req, - Response: new(resource.ResourceWrapper), - } - query, err := sqltemplate.Execute(sqlResourceList, readReq) - if err != nil { - return nil, fmt.Errorf("execute SQL template to list resources: %w", err) - } - - rows, err := b.sqlDB.QueryContext(ctx, query, readReq.GetArgs()...) - if err != nil { - return nil, fmt.Errorf("list resources: %w", err) - } + return b.listLatest(ctx, req) +} +// listLatest fetches the resources from the resource table. +func (b *backend) listLatest(ctx context.Context, req *resource.ListRequest) (*resource.ListResponse, error) { out := &resource.ListResponse{ - Items: make([]*resource.ResourceWrapper, req.Limit), - ResourceVersion: 0, // TODO + Items: []*resource.ResourceWrapper{}, // TODO: we could pre-allocate the capacity if we estimate the number of items + ResourceVersion: 0, + } - } - for i := 1; rows.Next(); i++ { - if ctx.Err() != nil { - return nil, ctx.Err() + err := b.sqlDB.WithTx(ctx, ReadCommittedRO, func(ctx context.Context, tx db.Tx) error { + var err error + + // TODO: Here the lastest RV might be lower than the actual latest RV + // because delete events are not included in the resource table. + out.ResourceVersion, err = fetchLatestRV(ctx, tx) + if err != nil { + return err } - if err := rows.Scan(readReq.GetScanDest()...); err != nil { - return nil, fmt.Errorf("scan row #%d: %w", i, err) + + // Fetch one extra row for Limit + if req.Limit > 0 { + req.Limit++ } - rw := *readReq.Response - out.Items = append(out.Items, &rw) - } - return out, nil + readReq := sqlResourceListRequest{ + SQLTemplate: sqltemplate.New(b.sqlDialect), + Request: req, + Response: new(resource.ResourceWrapper), + } + query, err := sqltemplate.Execute(sqlResourceList, readReq) + if err != nil { + return fmt.Errorf("execute SQL template to list resources: %w", err) + } + rows, err := tx.QueryContext(ctx, query, readReq.GetArgs()...) + if err != nil { + return fmt.Errorf("list resources: %w", err) + } + for i := int64(1); rows.Next(); i++ { + if ctx.Err() != nil { + return ctx.Err() + } + if err := rows.Scan(readReq.GetScanDest()...); err != nil { + return fmt.Errorf("scan row #%d: %w", i, err) + } + rw := *readReq.Response + + if req.Limit > 0 && i >= req.Limit { + continueToken := &ContinueToken{ResourceVersion: out.ResourceVersion, StartOffset: i - 1} + out.NextPageToken = continueToken.String() + break + } + out.Items = append(out.Items, &rw) + } + + return nil + }) + + return out, err } func (b *backend) WatchWriteEvents(ctx context.Context) (<-chan *resource.WrittenEvent, error) { @@ -424,19 +449,21 @@ func (b *backend) poller(ctx context.Context, since int64, stream chan<- *resour } // fetchLatestRV returns the current maxium RV in the resource table -func fetchLatestRV(ctx context.Context, db db.DB) (int64, error) { +func fetchLatestRV(ctx context.Context, db db.ContextExecer) (int64, error) { // Fetch the lastest RV rows, err := db.QueryContext(ctx, `SELECT COALESCE(max("resource_version"), 0) FROM "resource";`) if err != nil { return 0, fmt.Errorf("fetch latest rv: %w", err) } - since := int64(0) if rows.Next() { - if err := rows.Scan(&since); err != nil { + rv := new(int64) + if err := rows.Scan(&rv); err != nil { return 0, fmt.Errorf("scan since resource version: %w", err) } + return *rv, nil + } - return since, nil + return 0, fmt.Errorf("no rows") } func (b *backend) poll(ctx context.Context, since int64, stream chan<- *resource.WrittenEvent) (int64, error) { diff --git a/pkg/storage/unified/sql/backend_test.go b/pkg/storage/unified/sql/backend_test.go index 6a829c3fe17d..3c89d26e32a6 100644 --- a/pkg/storage/unified/sql/backend_test.go +++ b/pkg/storage/unified/sql/backend_test.go @@ -115,6 +115,7 @@ func TestBackendHappyPath(t *testing.T) { assert.Len(t, resp.Items, 2) assert.Equal(t, "updated value", string(resp.Items[0].Value)) assert.Equal(t, "initial value 3", string(resp.Items[1].Value)) + assert.Equal(t, int64(4), resp.ResourceVersion) }) t.Run("Watch events", func(t *testing.T) { @@ -188,3 +189,52 @@ func TestBackendWatchWriteEventsFromLastest(t *testing.T) { assert.NoError(t, err) assert.Equal(t, "item2", (<-stream).Key.Name) } + +func TestBackendPrepareList(t *testing.T) { + ctx := context.Background() + dbstore := db.InitTestDB(t) + + rdb, err := dbimpl.ProvideResourceDB(dbstore, setting.NewCfg(), featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorage), nil) + assert.NoError(t, err) + store, err := NewBackendStore(backendOptions{ + DB: rdb, + }) + + assert.NoError(t, err) + assert.NotNil(t, store) + + // Create a few resources before initing the watch + for i := 1; i <= 10; i++ { + rv, err := store.WriteEvent(ctx, resource.WriteEvent{ + Type: resource.WatchEvent_ADDED, + Value: []byte("initial value " + strconv.Itoa(i)), + Key: &resource.ResourceKey{ + Namespace: "namespace", + Group: "group", + Resource: "resource", + Name: "item" + strconv.Itoa(i), + }, + }) + assert.NoError(t, err) + assert.Equal(t, int64(i), rv) + } + + t.Run("fetch all latest", func(t *testing.T) { + res, err := store.PrepareList(ctx, &resource.ListRequest{}) + assert.NoError(t, err) + assert.Len(t, res.Items, 10) + assert.Empty(t, res.NextPageToken) + }) + + t.Run("fetch first page", func(t *testing.T) { + res, err := store.PrepareList(ctx, &resource.ListRequest{ + Limit: 5, + }) + assert.NoError(t, err) + assert.Len(t, res.Items, 5) + continueToken, err := GetContinueToken(res.NextPageToken) + assert.NoError(t, err) + assert.Equal(t, int64(10), continueToken.ResourceVersion) + assert.Equal(t, int64(5), continueToken.StartOffset) + }) +} diff --git a/pkg/storage/unified/sql/continue.go b/pkg/storage/unified/sql/continue.go new file mode 100644 index 000000000000..9e92e8eedd0b --- /dev/null +++ b/pkg/storage/unified/sql/continue.go @@ -0,0 +1,32 @@ +package sql + +import ( + "encoding/base64" + "encoding/json" + "fmt" +) + +type ContinueToken struct { + StartOffset int64 `json:"o"` + ResourceVersion int64 `json:"v"` +} + +func (c *ContinueToken) String() string { + b, _ := json.Marshal(c) + return base64.StdEncoding.EncodeToString(b) +} + +func GetContinueToken(token string) (*ContinueToken, error) { + continueVal, err := base64.StdEncoding.DecodeString(token) + if err != nil { + return nil, fmt.Errorf("error decoding continue token") + } + + t := &ContinueToken{} + err = json.Unmarshal(continueVal, t) + if err != nil { + return nil, err + } + + return t, nil +} diff --git a/pkg/storage/unified/sql/data/resource_list.sql b/pkg/storage/unified/sql/data/resource_list.sql index 6d2d6f617de3..7b24980e6a90 100644 --- a/pkg/storage/unified/sql/data/resource_list.sql +++ b/pkg/storage/unified/sql/data/resource_list.sql @@ -17,4 +17,8 @@ SELECT AND {{ .Ident "name" }} = {{ .Arg .Request.Options.Key.Name }} {{ end }} {{ end }} + ORDER BY {{ .Ident "resource_version" }} DESC + {{ if (gt .Request.Limit 0) }} + LIMIT {{ .Arg .Request.Limit }} + {{ end }} ; diff --git a/pkg/storage/unified/sql/queries_test.go b/pkg/storage/unified/sql/queries_test.go index 8df44dc84b65..da4cf1eb3b39 100644 --- a/pkg/storage/unified/sql/queries_test.go +++ b/pkg/storage/unified/sql/queries_test.go @@ -168,6 +168,7 @@ func TestQueries(t *testing.T) { Data: &sqlResourceListRequest{ SQLTemplate: new(sqltemplate.SQLTemplate), Request: &resource.ListRequest{ + Limit: 10, Options: &resource.ListOptions{ Key: &resource.ResourceKey{ Namespace: "ns", diff --git a/pkg/storage/unified/sql/testdata/resource_list_mysql_sqlite.sql b/pkg/storage/unified/sql/testdata/resource_list_mysql_sqlite.sql index f9aa439ca3f0..0b2edae1e314 100644 --- a/pkg/storage/unified/sql/testdata/resource_list_mysql_sqlite.sql +++ b/pkg/storage/unified/sql/testdata/resource_list_mysql_sqlite.sql @@ -1,3 +1,6 @@ SELECT "resource_version", "value" FROM "resource" - WHERE 1 = 1 AND "namespace" = ?; + WHERE 1 = 1 AND "namespace" = ? + ORDER BY "resource_version" DESC + LIMIT ? +;