diff --git a/components/usage/pkg/apiv1/report-generator.go b/components/usage/pkg/apiv1/report-generator.go index 02ecadb7a3f33c..60a1cd5f1efd9e 100644 --- a/components/usage/pkg/apiv1/report-generator.go +++ b/components/usage/pkg/apiv1/report-generator.go @@ -7,9 +7,10 @@ package apiv1 import ( "context" "fmt" - "github.com/gitpod-io/gitpod/usage/pkg/contentservice" "time" + "github.com/gitpod-io/gitpod/usage/pkg/contentservice" + "github.com/gitpod-io/gitpod/common-go/log" "github.com/gitpod-io/gitpod/usage/pkg/db" "google.golang.org/grpc/codes" @@ -76,15 +77,15 @@ func validateInstances(instances []db.WorkspaceInstanceForUsage) (valid []db.Wor instance := i // Each instance must have a start time, without it, we do not have a baseline for usage computation. - if !instance.CreationTime.IsSet() { + if !instance.StartedTime.IsSet() { invalid = append(invalid, contentservice.InvalidSession{ - Reason: "missing creation time", + Reason: "missing started time", Session: instance, }) continue } - start := instance.CreationTime.Time() + start := instance.StartedTime.Time() // Currently running instances do not have a stopped time set, so we ignore these. if instance.StoppingTime.IsSet() { diff --git a/components/usage/pkg/apiv1/usage_test.go b/components/usage/pkg/apiv1/usage_test.go index e219c54b281334..938f84d50c6f4f 100644 --- a/components/usage/pkg/apiv1/usage_test.go +++ b/components/usage/pkg/apiv1/usage_test.go @@ -7,11 +7,12 @@ package apiv1 import ( "context" "database/sql" - "github.com/gitpod-io/gitpod/usage/pkg/contentservice" "reflect" "testing" "time" + "github.com/gitpod-io/gitpod/usage/pkg/contentservice" + "github.com/gitpod-io/gitpod/common-go/baseserver" v1 "github.com/gitpod-io/gitpod/usage-api/v1" "github.com/gitpod-io/gitpod/usage/pkg/db" @@ -104,7 +105,6 @@ func TestUsageService_ListBilledUsage(t *testing.T) { start := time.Date(2022, 07, 1, 13, 0, 0, 0, time.UTC) attrID := db.NewTeamAttributionID(uuid.New().String()) var instances []db.WorkspaceInstanceUsage - var instanceIDs []string for i := 0; i < 3; i++ { instance := dbtest.NewWorkspaceInstanceUsage(t, db.WorkspaceInstanceUsage{ AttributionID: attrID, @@ -115,8 +115,6 @@ func TestUsageService_ListBilledUsage(t *testing.T) { }, }) instances = append(instances, instance) - - instanceIDs = append(instanceIDs, instance.InstanceID.String()) } return Scenario{ @@ -395,25 +393,20 @@ func TestReportGenerator_GenerateUsageReport(t *testing.T) { dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), UsageAttributionID: db.NewTeamAttributionID(teamID.String()), - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 01, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 1, 0, 0, time.UTC)), }), // Still running dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), UsageAttributionID: db.NewTeamAttributionID(teamID.String()), - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 01, 00, 00, time.UTC)), }), // No creation time, invalid record dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), UsageAttributionID: db.NewTeamAttributionID(teamID.String()), - StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 1, 0, 0, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 1, 0, 0, time.UTC)), }), } diff --git a/components/usage/pkg/db/dbtest/usage.go b/components/usage/pkg/db/dbtest/usage.go index 2e122af35cf165..a34712b0098249 100644 --- a/components/usage/pkg/db/dbtest/usage.go +++ b/components/usage/pkg/db/dbtest/usage.go @@ -5,6 +5,7 @@ package dbtest import ( + "context" "testing" "github.com/gitpod-io/gitpod/usage/pkg/db" @@ -67,8 +68,7 @@ func CreateUsageRecords(t *testing.T, conn *gorm.DB, entries ...db.Usage) []db.U ids = append(ids, record.ID.String()) } - require.NoError(t, conn.CreateInBatches(&records, 1000).Error) - + require.NoError(t, db.InsertUsage(context.Background(), conn, entries...)) t.Cleanup(func() { require.NoError(t, conn.Where(ids).Delete(&db.Usage{}).Error) }) diff --git a/components/usage/pkg/db/dbtest/workspace_instance.go b/components/usage/pkg/db/dbtest/workspace_instance.go index 5ddc244c8509ed..a2f4671016af87 100644 --- a/components/usage/pkg/db/dbtest/workspace_instance.go +++ b/components/usage/pkg/db/dbtest/workspace_instance.go @@ -6,12 +6,13 @@ package dbtest import ( "database/sql" + "testing" + "time" + "github.com/gitpod-io/gitpod/usage/pkg/db" "github.com/google/uuid" "github.com/stretchr/testify/require" "gorm.io/gorm" - "testing" - "time" ) var ( @@ -34,6 +35,8 @@ func NewWorkspaceInstance(t *testing.T, instance db.WorkspaceInstance) db.Worksp creationTime := db.VarcharTime{} if instance.CreationTime.IsSet() { creationTime = instance.CreationTime + } else if instance.StartedTime.IsSet() { + creationTime = instance.StartedTime } startedTime := db.VarcharTime{} @@ -49,6 +52,8 @@ func NewWorkspaceInstance(t *testing.T, instance db.WorkspaceInstance) db.Worksp stoppedTime := db.VarcharTime{} if instance.StoppedTime.IsSet() { stoppedTime = instance.StoppedTime + } else if instance.StoppingTime.IsSet() { + creationTime = instance.StoppingTime } stoppingTime := db.VarcharTime{} diff --git a/components/usage/pkg/db/usage.go b/components/usage/pkg/db/usage.go index 4e72835a609a43..ee470dd2a7e750 100644 --- a/components/usage/pkg/db/usage.go +++ b/components/usage/pkg/db/usage.go @@ -11,6 +11,7 @@ import ( "github.com/google/uuid" "gorm.io/datatypes" "gorm.io/gorm" + "gorm.io/gorm/clause" ) type Usage struct { @@ -34,18 +35,47 @@ func (u *Usage) TableName() string { return "d_b_usage" } -func FindUsage(ctx context.Context, conn *gorm.DB, attributionId AttributionID, from, to VarcharTime, offset int64, limit int64) ([]Usage, error) { - db := conn.WithContext(ctx) +func InsertUsage(ctx context.Context, conn *gorm.DB, records ...Usage) error { + return conn.WithContext(ctx). + Clauses(clause.OnConflict{DoNothing: true}). + CreateInBatches(records, 1000).Error +} + +func UpdateUsage(ctx context.Context, conn *gorm.DB, record Usage) error { + return conn.WithContext(ctx).Save(record).Error +} +func FindAllDraftUsage(ctx context.Context, conn *gorm.DB) ([]Usage, error) { var usageRecords []Usage - result := db. - WithContext(ctx). + var usageRecordsBatch []Usage + + result := conn.WithContext(ctx). + Where("draft = TRUE"). + Order("effectiveTime DESC"). + FindInBatches(&usageRecordsBatch, 1000, func(_ *gorm.DB, _ int) error { + usageRecords = append(usageRecords, usageRecordsBatch...) + return nil + }) + if result.Error != nil { + return nil, fmt.Errorf("failed to get usage records: %s", result.Error) + } + return usageRecords, nil +} + +func FindUsage(ctx context.Context, conn *gorm.DB, attributionId AttributionID, from, to VarcharTime, offset int64, limit int64) ([]Usage, error) { + var usageRecords []Usage + var usageRecordsBatch []Usage + + result := conn.WithContext(ctx). Where("attributionId = ?", attributionId). Where("? <= effectiveTime AND effectiveTime < ?", from.String(), to.String()). Order("effectiveTime DESC"). Offset(int(offset)). Limit(int(limit)). - Find(&usageRecords) + FindInBatches(&usageRecordsBatch, 1000, func(_ *gorm.DB, _ int) error { + usageRecords = append(usageRecords, usageRecordsBatch...) + return nil + }) if result.Error != nil { return nil, fmt.Errorf("failed to get usage records: %s", result.Error) } diff --git a/components/usage/pkg/db/usage_test.go b/components/usage/pkg/db/usage_test.go index 270f576a041db6..6e6ad7071a4d48 100644 --- a/components/usage/pkg/db/usage_test.go +++ b/components/usage/pkg/db/usage_test.go @@ -47,3 +47,93 @@ func TestFindUsageInRange(t *testing.T) { require.Equal(t, 1, len(listResult)) require.Equal(t, []db.Usage{entryInside}, listResult) } + +func TestInsertUsageRecords(t *testing.T) { + conn := dbtest.ConnectForTests(t) + + attributionID := db.NewTeamAttributionID(uuid.New().String()) + start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC) + + usage := dbtest.NewUsage(t, db.Usage{ + AttributionID: attributionID, + EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)), + Draft: true, + }) + + dbtest.CreateUsageRecords(t, conn, usage) + updatedDesc := "Updated Description" + usage.Description = updatedDesc + + require.NoError(t, db.InsertUsage(context.Background(), conn, usage)) + + drafts, err := db.FindAllDraftUsage(context.Background(), conn) + require.NoError(t, err) + require.Equal(t, 1, len(drafts)) + require.NotEqual(t, updatedDesc, drafts[0].Description) +} + +func TestUpdateUsageRecords(t *testing.T) { + conn := dbtest.ConnectForTests(t) + + attributionID := db.NewTeamAttributionID(uuid.New().String()) + start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC) + + usage := dbtest.NewUsage(t, db.Usage{ + AttributionID: attributionID, + EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)), + Draft: true, + }) + + dbtest.CreateUsageRecords(t, conn, usage) + updatedDesc := "Updated Description" + usage.Description = updatedDesc + + require.NoError(t, db.UpdateUsage(context.Background(), conn, usage)) + + drafts, err := db.FindAllDraftUsage(context.Background(), conn) + require.NoError(t, err) + require.Equal(t, 1, len(drafts)) + require.Equal(t, updatedDesc, drafts[0].Description) +} + +func TestFindAllDraftUsage(t *testing.T) { + conn := dbtest.ConnectForTests(t) + + attributionID := db.NewTeamAttributionID(uuid.New().String()) + start := time.Date(2022, 7, 1, 0, 0, 0, 0, time.UTC) + + usage1 := dbtest.NewUsage(t, db.Usage{ + AttributionID: attributionID, + EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)), + Draft: true, + }) + usage2 := dbtest.NewUsage(t, db.Usage{ + AttributionID: attributionID, + EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)), + Draft: true, + }) + usage3 := dbtest.NewUsage(t, db.Usage{ + AttributionID: attributionID, + EffectiveTime: db.NewVarcharTime(start.Add(2 * time.Hour)), + Draft: false, + }) + + dbtest.CreateUsageRecords(t, conn, usage1, usage2, usage3) + drafts, err := db.FindAllDraftUsage(context.Background(), conn) + require.NoError(t, err) + require.Equal(t, 2, len(drafts)) + for _, usage := range drafts { + require.True(t, usage.Draft) + } + + // let's finalize one record + usage2.Draft = false + require.NoError(t, db.UpdateUsage(context.Background(), conn, usage2)) + + drafts, err = db.FindAllDraftUsage(context.Background(), conn) + require.NoError(t, err) + require.Equal(t, 1, len(drafts)) + for _, usage := range drafts { + require.True(t, usage.Draft) + } +} diff --git a/components/usage/pkg/db/workspace_instance.go b/components/usage/pkg/db/workspace_instance.go index 7f3efea55086cb..cb3f1f6cbd9e81 100644 --- a/components/usage/pkg/db/workspace_instance.go +++ b/components/usage/pkg/db/workspace_instance.go @@ -51,6 +51,75 @@ func (i *WorkspaceInstance) TableName() string { return "d_b_workspace_instance" } +// FindStoppedWorkspaceInstancesInRange finds WorkspaceInstanceForUsage that have been stopped between from (inclusive) and to (exclusive). +func FindStoppedWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to time.Time) ([]WorkspaceInstanceForUsage, error) { + var instances []WorkspaceInstanceForUsage + var instancesInBatch []WorkspaceInstanceForUsage + + tx := queryWorkspaceInstanceForUsage(ctx, conn). + Where("wsi.stoppingTime >= ?", TimeToISO8601(from)). + Where("wsi.stoppingTime < ?", TimeToISO8601(to)). + Where("wsi.stoppingTime != ?", ""). + Where("wsi.usageAttributionId != ?", ""). + FindInBatches(&instancesInBatch, 1000, func(_ *gorm.DB, _ int) error { + instances = append(instances, instancesInBatch...) + return nil + }) + if tx.Error != nil { + return nil, fmt.Errorf("failed to find workspace instances: %w", tx.Error) + } + + return instances, nil +} + +// FindRunningWorkspaceInstances finds WorkspaceInstanceForUsage that are running at the point in time the querty is executed. +func FindRunningWorkspaceInstances(ctx context.Context, conn *gorm.DB) ([]WorkspaceInstanceForUsage, error) { + var instances []WorkspaceInstanceForUsage + var instancesInBatch []WorkspaceInstanceForUsage + + tx := queryWorkspaceInstanceForUsage(ctx, conn). + Where("wsi.stoppingTime = ?", ""). + Where("wsi.usageAttributionId != ?", ""). + FindInBatches(&instancesInBatch, 1000, func(_ *gorm.DB, _ int) error { + instances = append(instances, instancesInBatch...) + return nil + }) + if tx.Error != nil { + return nil, fmt.Errorf("failed to find running workspace instances: %w", tx.Error) + } + + return instances, nil +} + +// FindWorkspaceInstancesByIds finds WorkspaceInstanceForUsage by Id. +func FindWorkspaceInstancesByIds(ctx context.Context, conn *gorm.DB, workspaceInstanceIds []uuid.UUID) ([]WorkspaceInstanceForUsage, error) { + var instances []WorkspaceInstanceForUsage + var instancesInBatch []WorkspaceInstanceForUsage + var idChunks [][]uuid.UUID + chunkSize, totalSize := 1000, len(workspaceInstanceIds) + // explicit batching to reduce the lengths of the 'in'-part in the SELECT statement below + for i := 0; i < totalSize; i += chunkSize { + end := i + chunkSize + if end > totalSize { + end = totalSize + } + idChunks = append(idChunks, workspaceInstanceIds[i:end]) + } + + for _, idChunk := range idChunks { + err := queryWorkspaceInstanceForUsage(ctx, conn). + Where("wsi.id in ?", idChunk). + Where("wsi.usageAttributionId != ?", ""). + Find(&instancesInBatch).Error + if err != nil { + return nil, fmt.Errorf("failed to find workspace instances by id: %w", err) + } + instances = append(instances, instancesInBatch...) + } + + return instances, nil +} + // ListWorkspaceInstancesInRange lists WorkspaceInstances between from (inclusive) and to (exclusive). // This results in all instances which have existed in the specified period, regardless of their current status, this includes: // - terminated @@ -61,26 +130,11 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to var instances []WorkspaceInstanceForUsage var instancesInBatch []WorkspaceInstanceForUsage - tx := conn.WithContext(ctx). - Table(fmt.Sprintf("%s as wsi", (&WorkspaceInstance{}).TableName())). - Select("wsi.id as id, "+ - "ws.projectId as projectId, "+ - "ws.type as workspaceType, "+ - "wsi.workspaceClass as workspaceClass, "+ - "wsi.usageAttributionId as usageAttributionId, "+ - "wsi.creationTime as creationTime, "+ - "wsi.startedTime as startedTime, "+ - "wsi.stoppingTime as stoppingTime, "+ - "wsi.stoppedTime as stoppedTime, "+ - "ws.ownerId as ownerId, "+ - "ws.id as workspaceId", - ). - Joins(fmt.Sprintf("LEFT JOIN %s AS ws ON wsi.workspaceId = ws.id", (&Workspace{}).TableName())). + tx := queryWorkspaceInstanceForUsage(ctx, conn). Where( conn.Where("wsi.stoppingTime >= ?", TimeToISO8601(from)).Or("wsi.stoppingTime = ?", ""), ). - Where("wsi.creationTime < ?", TimeToISO8601(to)). - Where("wsi.startedTime != ?", ""). + Where("wsi.startedTime < ?", TimeToISO8601(to)). Where("wsi.usageAttributionId != ?", ""). FindInBatches(&instancesInBatch, 1000, func(_ *gorm.DB, _ int) error { instances = append(instances, instancesInBatch...) @@ -93,6 +147,24 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to return instances, nil } +func queryWorkspaceInstanceForUsage(ctx context.Context, conn *gorm.DB) *gorm.DB { + return conn.WithContext(ctx). + Table(fmt.Sprintf("%s as wsi", (&WorkspaceInstance{}).TableName())). + Select("wsi.id as id, " + + "ws.projectId as projectId, " + + "ws.type as workspaceType, " + + "wsi.workspaceClass as workspaceClass, " + + "wsi.usageAttributionId as usageAttributionId, " + + "wsi.creationTime as creationTime, " + + "wsi.startedTime as startedTime, " + + "wsi.stoppingTime as stoppingTime, " + + "wsi.stoppedTime as stoppedTime, " + + "ws.ownerId as ownerId, " + + "ws.id as workspaceId", + ). + Joins(fmt.Sprintf("LEFT JOIN %s AS ws ON wsi.workspaceId = ws.id", (&Workspace{}).TableName())) +} + const ( AttributionEntity_User = "user" AttributionEntity_Team = "team" diff --git a/components/usage/pkg/db/workspace_instance_test.go b/components/usage/pkg/db/workspace_instance_test.go index 2d3d4fb4c0369d..b294a782dd4479 100644 --- a/components/usage/pkg/db/workspace_instance_test.go +++ b/components/usage/pkg/db/workspace_instance_test.go @@ -22,7 +22,7 @@ var ( startOfJune = time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC) ) -func TestListWorkspaceInstancesInRange(t *testing.T) { +func TestFindStoppedWorkspaceInstancesInRange(t *testing.T) { conn := dbtest.ConnectForTests(t) workspace := dbtest.CreateWorkspaces(t, conn, dbtest.NewWorkspace(t, db.Workspace{}))[0] @@ -31,81 +31,58 @@ func TestListWorkspaceInstancesInRange(t *testing.T) { // In the middle of May dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ WorkspaceID: workspace.ID, - 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)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), }), // Start of May dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 1, 1, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 1, 00, 00, 00, time.UTC)), }), // End of May dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 59, 59, 999999, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 59, 59, 999999, time.UTC)), }), // Started in April, but continued into May dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 04, 30, 23, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 04, 30, 23, 00, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 0, 0, 0, time.UTC)), + }), + } + invalid := []db.WorkspaceInstance{ + // Started in April, no stop time, still running + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), }), // Started in May, but continued into June dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), }), // Started in April, but continued into June (ran for all of May) dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), }), - // Stopped in May, no creation time, should be retrieved but this is a poor data quality record. - dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - ID: uuid.New(), - WorkspaceID: workspace.ID, - StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 1, 0, 0, 0, time.UTC)), - StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 1, 0, 0, 0, time.UTC)), - }), - // Started in April, no stop time, still running - dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - ID: uuid.New(), - WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), - StartedTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), - }), - } - invalid := []db.WorkspaceInstance{ // Start of June dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ ID: uuid.New(), WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 06, 1, 00, 00, 00, 00, time.UTC)), StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 00, 00, 00, 00, time.UTC)), StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), }), } @@ -115,7 +92,7 @@ func TestListWorkspaceInstancesInRange(t *testing.T) { dbtest.CreateWorkspaceInstances(t, conn, all...) - retrieved, err := db.ListWorkspaceInstancesInRange(context.Background(), conn, startOfMay, startOfJune) + retrieved, err := db.FindStoppedWorkspaceInstancesInRange(context.Background(), conn, startOfMay, startOfJune) require.NoError(t, err) require.Len(t, retrieved, len(valid)) @@ -129,10 +106,8 @@ func TestListWorkspaceInstancesInRange_Fields(t *testing.T) { workspace := dbtest.CreateWorkspaces(t, conn, dbtest.NewWorkspace(t, db.Workspace{}))[0] instance := dbtest.CreateWorkspaceInstances(t, conn, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ WorkspaceID: workspace.ID, - 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)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), }))[0] retrieved, err := db.ListWorkspaceInstancesInRange(context.Background(), conn, startOfMay, startOfJune) @@ -165,10 +140,8 @@ func TestListWorkspaceInstancesInRange_Fields(t *testing.T) { }))[0] instance := dbtest.CreateWorkspaceInstances(t, conn, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ WorkspaceID: workspace.ID, - 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)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), }))[0] retrieved, err := db.ListWorkspaceInstancesInRange(context.Background(), conn, startOfMay, startOfJune) @@ -204,10 +177,8 @@ func TestListWorkspaceInstancesInRange_InBatches(t *testing.T) { 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)), StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), })) } @@ -240,3 +211,159 @@ func TestAttributionID_Values(t *testing.T) { }) } } + +func TestListWorkspaceInstancesInRange(t *testing.T) { + conn := dbtest.ConnectForTests(t) + + workspace := dbtest.CreateWorkspaces(t, conn, dbtest.NewWorkspace(t, db.Workspace{}))[0] + + valid := []db.WorkspaceInstance{ + // In the middle of May + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 12, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), + }), + // Start of May + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 1, 1, 00, 00, 00, time.UTC)), + }), + // End of May + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 59, 59, 999999, time.UTC)), + }), + // Started in April, but continued into May + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 04, 30, 23, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 0, 0, 0, time.UTC)), + }), + // Started in May, but continued into June + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), + }), + // Started in April, but continued into June (ran for all of May) + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), + }), + // Still Running + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 00, 00, 00, 00, time.UTC)), + }), + } + invalid := []db.WorkspaceInstance{ + // Started in April, no stop time, still running + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 04, 31, 23, 00, 00, 00, time.UTC)), + }), + } + + var all []db.WorkspaceInstance + all = append(all, valid...) + all = append(all, invalid...) + + dbtest.CreateWorkspaceInstances(t, conn, all...) + + retrieved, err := db.ListWorkspaceInstancesInRange(context.Background(), conn, startOfMay, startOfJune) + require.NoError(t, err) + + ids := []uuid.UUID{} + for _, ws := range retrieved { + ids = append(ids, ws.ID) + } + + require.Len(t, retrieved, len(valid)) +} + +func TestFindRunningWorkspace(t *testing.T) { + conn := dbtest.ConnectForTests(t) + + workspace := dbtest.CreateWorkspaces(t, conn, dbtest.NewWorkspace(t, db.Workspace{}))[0] + + all := []db.WorkspaceInstance{ + // one stopped instance + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 12, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), + }), + // Two running instances + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)), + }), + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 04, 30, 23, 00, 00, 00, time.UTC)), + }), + } + + dbtest.CreateWorkspaceInstances(t, conn, all...) + + retrieved, err := db.FindRunningWorkspaceInstances(context.Background(), conn) + require.NoError(t, err) + + require.Equal(t, 2, len(retrieved)) + for _, ws := range retrieved { + require.False(t, ws.StoppingTime.IsSet()) + } + +} + +func TestFindWorkspacesByInstanceId(t *testing.T) { + conn := dbtest.ConnectForTests(t) + + workspace := dbtest.CreateWorkspaces(t, conn, dbtest.NewWorkspace(t, db.Workspace{}))[0] + + all := []db.WorkspaceInstance{ + // one stopped instance + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 15, 12, 00, 00, 00, time.UTC)), + StoppingTime: db.NewVarcharTime(time.Date(2022, 05, 15, 13, 00, 00, 00, time.UTC)), + }), + // Two running instances + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC)), + }), + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + WorkspaceID: workspace.ID, + StartedTime: db.NewVarcharTime(time.Date(2022, 04, 30, 23, 00, 00, 00, time.UTC)), + }), + } + + twoIds := []uuid.UUID{all[0].ID, all[1].ID} + + dbtest.CreateWorkspaceInstances(t, conn, all...) + + retrieved, err := db.FindWorkspaceInstancesByIds(context.Background(), conn, twoIds) + require.NoError(t, err) + + require.Equal(t, 2, len(retrieved)) + for _, ws := range retrieved { + require.NotEqual(t, all[2].ID, ws.ID) + } + +}