Skip to content

[usage] Add database functions for new usage control loop #12643

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 5, 2022
Merged

[usage] Add database functions for new usage control loop #12643

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f6d86c6
[usage] store and fetch usage entries
svenefftinge Sep 5, 2022
4bdf2c0
[usage] more workspace instance fetch methods
svenefftinge Sep 5, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions components/usage/pkg/apiv1/report-generator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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() {
Expand Down
11 changes: 2 additions & 9 deletions components/usage/pkg/apiv1/usage_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand All @@ -115,8 +115,6 @@ func TestUsageService_ListBilledUsage(t *testing.T) {
},
})
instances = append(instances, instance)

instanceIDs = append(instanceIDs, instance.InstanceID.String())
}

return Scenario{
Expand Down Expand Up @@ -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)),
}),
}

Expand Down
4 changes: 2 additions & 2 deletions components/usage/pkg/db/dbtest/usage.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
package dbtest

import (
"context"
"testing"

"github.com/gitpod-io/gitpod/usage/pkg/db"
Expand Down Expand Up @@ -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)
})
Expand Down
9 changes: 7 additions & 2 deletions components/usage/pkg/db/dbtest/workspace_instance.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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{}
Expand All @@ -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{}
Expand Down
40 changes: 35 additions & 5 deletions components/usage/pkg/db/usage.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"github.com/google/uuid"
"gorm.io/datatypes"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)

type Usage struct {
Expand All @@ -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)
}
Expand Down
90 changes: 90 additions & 0 deletions components/usage/pkg/db/usage_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

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)
}
}
Loading