diff --git a/components/usage/go.mod b/components/usage/go.mod index 648af860e81898..444d4d6458588f 100644 --- a/components/usage/go.mod +++ b/components/usage/go.mod @@ -63,6 +63,7 @@ require ( github.com/sirupsen/logrus v1.8.1 github.com/spf13/cobra v1.4.0 github.com/stretchr/testify v1.7.0 + github.com/stripe/stripe-go/v72 v72.114.0 gorm.io/datatypes v1.0.6 gorm.io/driver/mysql v1.3.3 gorm.io/gorm v1.23.5 @@ -88,7 +89,6 @@ require ( github.com/prometheus/common v0.32.1 // indirect github.com/prometheus/procfs v0.7.3 // indirect github.com/spf13/pflag v1.0.5 // indirect - github.com/stripe/stripe-go/v72 v72.114.0 // indirect golang.org/x/net v0.0.0-20211209124913-491a49abca63 // indirect golang.org/x/sync v0.0.0-20210220032951-036812b2e83c // indirect golang.org/x/sys v0.0.0-20220114195835-da31bd327af9 // indirect diff --git a/components/usage/pkg/db/dbtest/workspace.go b/components/usage/pkg/db/dbtest/workspace.go index 9c10251921bd47..cb6b998dbce271 100644 --- a/components/usage/pkg/db/dbtest/workspace.go +++ b/components/usage/pkg/db/dbtest/workspace.go @@ -23,7 +23,7 @@ const ( func NewWorkspace(t *testing.T, workspace db.Workspace) db.Workspace { t.Helper() - id := generateWorkspaceID() + id := GenerateWorkspaceID() if workspace.ID != "" { id = workspace.ID } @@ -63,7 +63,7 @@ func NewWorkspace(t *testing.T, workspace db.Workspace) db.Workspace { } } -func generateWorkspaceID() string { +func GenerateWorkspaceID() string { return fmt.Sprintf("gitpodio-gitpod-%s", randSeq(11)) } diff --git a/components/usage/pkg/db/workspace.go b/components/usage/pkg/db/workspace.go index 1cc4fd8b91dff3..827fb2d2a0847c 100644 --- a/components/usage/pkg/db/workspace.go +++ b/components/usage/pkg/db/workspace.go @@ -11,6 +11,7 @@ import ( "github.com/google/uuid" "gorm.io/datatypes" "gorm.io/gorm" + "math" "time" ) @@ -59,15 +60,33 @@ const ( WorkspaceType_Regular WorkspaceType = "regular" ) +const maxListBatchSize = 65535 // 2^16 - 1 + func ListWorkspacesByID(ctx context.Context, conn *gorm.DB, ids []string) ([]Workspace, error) { if len(ids) == 0 { return nil, nil } var workspaces []Workspace - tx := conn.WithContext(ctx).Where(ids).Find(&workspaces) - if tx.Error != nil { - return nil, fmt.Errorf("failed to list workspaces by id: %w", tx.Error) + + items := len(ids) + batches := int(math.Ceil(float64(items) / maxListBatchSize)) + for i := 0; i < batches; i++ { + lower := i * maxListBatchSize + upper := (i + 1) * maxListBatchSize + + if upper > items { + upper = items + } + + batchIDs := ids[lower:upper] + var results []Workspace + tx := conn.WithContext(ctx).Where(batchIDs).Find(&results) + if tx.Error != nil { + return nil, fmt.Errorf("failed to list workspaces by id: %w", tx.Error) + } + + workspaces = append(workspaces, results...) } return workspaces, nil diff --git a/components/usage/pkg/db/workspace_test.go b/components/usage/pkg/db/workspace_test.go index dad4ba27dd638c..de0fc40f9edfaa 100644 --- a/components/usage/pkg/db/workspace_test.go +++ b/components/usage/pkg/db/workspace_test.go @@ -126,6 +126,11 @@ func TestListWorkspacesByID(t *testing.T) { QueryIDs: []string{workspaces[0].ID, workspaces[1].ID}, Expected: 2, }, + { + Name: "over 2^16 - 1 IDs requires batching", + QueryIDs: append([]string{workspaces[0].ID, workspaces[1].ID}, generateWorkspaceIDs(65535)...), + Expected: 2, + }, } { t.Run(scenario.Name, func(t *testing.T) { results, err := db.ListWorkspacesByID(context.Background(), conn, scenario.QueryIDs) @@ -135,3 +140,12 @@ func TestListWorkspacesByID(t *testing.T) { } } + +func generateWorkspaceIDs(count int) []string { + var ids []string + for i := 0; i < count; i++ { + ids = append(ids, dbtest.GenerateWorkspaceID()) + } + + return ids +}