From febbad0c8c7a870a49cba330105a9c7bb853231a Mon Sep 17 00:00:00 2001 From: Milan Pavlik Date: Tue, 21 Jun 2022 08:12:28 +0000 Subject: [PATCH] [usage] Batch list workspace instances in range --- .../usage/pkg/db/dbtest/workspace_instance.go | 2 +- components/usage/pkg/db/workspace_instance.go | 6 ++++- .../usage/pkg/db/workspace_instance_test.go | 26 +++++++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/components/usage/pkg/db/dbtest/workspace_instance.go b/components/usage/pkg/db/dbtest/workspace_instance.go index 251a414c2c7cef..72dc7604d63581 100644 --- a/components/usage/pkg/db/dbtest/workspace_instance.go +++ b/components/usage/pkg/db/dbtest/workspace_instance.go @@ -24,7 +24,7 @@ func NewWorkspaceInstance(t *testing.T, instance db.WorkspaceInstance) db.Worksp id = instance.ID } - workspaceID := generateWorkspaceID() + workspaceID := GenerateWorkspaceID() if instance.WorkspaceID != "" { workspaceID = instance.WorkspaceID } diff --git a/components/usage/pkg/db/workspace_instance.go b/components/usage/pkg/db/workspace_instance.go index 32c0252022ffa0..fdcf0fe7af49d5 100644 --- a/components/usage/pkg/db/workspace_instance.go +++ b/components/usage/pkg/db/workspace_instance.go @@ -70,13 +70,17 @@ func (i *WorkspaceInstance) TableName() string { // - instances which only just started in the period specified func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to time.Time) ([]WorkspaceInstance, error) { var instances []WorkspaceInstance + var instancesInBatch []WorkspaceInstance tx := conn.WithContext(ctx). Where( conn.Where("stoppedTime >= ?", TimeToISO8601(from)).Or("stoppedTime = ?", ""), ). Where("creationTime < ?", TimeToISO8601(to)). Where("startedTime != ?", ""). - Find(&instances) + FindInBatches(&instancesInBatch, 1000, func(_ *gorm.DB, _ int) error { + instances = append(instances, instancesInBatch...) + return nil + }) if tx.Error != nil { return nil, fmt.Errorf("failed to list workspace instances: %w", tx.Error) } diff --git a/components/usage/pkg/db/workspace_instance_test.go b/components/usage/pkg/db/workspace_instance_test.go index 614a057bd79b56..be9cf739880457 100644 --- a/components/usage/pkg/db/workspace_instance_test.go +++ b/components/usage/pkg/db/workspace_instance_test.go @@ -191,3 +191,29 @@ func TestListWorkspaceInstancesInRange(t *testing.T) { require.Len(t, retrieved, len(valid)) } + +func TestListWorkspaceInstancesInRange_InBatches(t *testing.T) { + conn := db.ConnectForTests(t) + + workspaceID := "gitpodio-gitpod-gyjr82jkfnd" + var instances []db.WorkspaceInstance + for i := 0; i < 1100; i++ { + instances = append(instances, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspaceID, + CreationTime: db.NewVarcharTime(time.Date(2022, 05, 15, 12, 00, 00, 00, time.UTC)), + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 12, 00, 00, 00, time.UTC)), + StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), + })) + + } + + tx := conn.CreateInBatches(&instances, 1000) + require.NoError(t, tx.Error) + + startOfMay := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC) + startOfJune := time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC) + results, err := db.ListWorkspaceInstancesInRange(context.Background(), conn, startOfMay, startOfJune) + require.NoError(t, err) + require.Len(t, results, len(instances)) +}