usage.go

// Copyright (c) 2022 Gitpod GmbH. All rights reserved.
// Licensed under the GNU Affero General Public License (AGPL).
// See License.AGPL.txt in the project root for license information.

package db

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"time"

	"github.com/google/uuid"
	"gorm.io/datatypes"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type UsageKind string

const (
	WorkspaceInstanceUsageKind UsageKind = "workspaceinstance"
	InvoiceUsageKind           UsageKind = "invoice"
	CreditNoteKind             UsageKind = "creditnote"
)

func NewCreditCents(n float64) CreditCents {
	inCents := n * 100
	return CreditCents(int64(math.Round(inCents)))
}

type CreditCents int64

func (cc CreditCents) ToCredits() float64 {
	return float64(cc) / 100
}

type Usage struct {
	ID                  uuid.UUID      `gorm:"primary_key;column:id;type:char;size:36;" json:"id"`
	AttributionID       AttributionID  `gorm:"column:attributionId;type:varchar;size:255;" json:"attributionId"`
	Description         string         `gorm:"column:description;type:varchar;size:255;" json:"description"`
	CreditCents         CreditCents    `gorm:"column:creditCents;type:bigint;" json:"creditCents"`
	EffectiveTime       VarcharTime    `gorm:"column:effectiveTime;type:varchar;size:255;" json:"effectiveTime"`
	Kind                UsageKind      `gorm:"column:kind;type:char;size:10;" json:"kind"`
	WorkspaceInstanceID *uuid.UUID     `gorm:"column:workspaceInstanceId;type:char;size:36;" json:"workspaceInstanceId"`
	Draft               bool           `gorm:"column:draft;type:boolean;" json:"draft"`
	Metadata            datatypes.JSON `gorm:"column:metadata;type:text;size:65535" json:"metadata"`
}

func (u *Usage) SetMetadataWithWorkspaceInstance(data WorkspaceInstanceUsageData) error {
	b, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to serialize workspace instance usage data into json: %w", err)
	}

	u.Metadata = b
	return nil
}

func (u *Usage) SetCreditNoteMetaData(data CreditNoteMetaData) error {
	b, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to serialize credit note meta data into json: %w", err)
	}

	u.Metadata = b
	return nil
}

func (u *Usage) GetMetadataAsWorkspaceInstanceData() (WorkspaceInstanceUsageData, error) {
	var data WorkspaceInstanceUsageData
	err := json.Unmarshal(u.Metadata, &data)
	if err != nil {
		return WorkspaceInstanceUsageData{}, fmt.Errorf("failed unmarshal metadata into wokrspace instance data: %w", err)
	}

	return data, nil
}

// WorkspaceInstanceUsageData represents the shape of metadata for usage entries of kind "workspaceinstance"
// the equivalent TypeScript definition is maintained in `components/gitpod-protocol/src/usage.ts“
type WorkspaceInstanceUsageData struct {
	WorkspaceId    string        `json:"workspaceId"`
	WorkspaceType  WorkspaceType `json:"workspaceType"`
	WorkspaceClass string        `json:"workspaceClass"`
	ContextURL     string        `json:"contextURL"`
	CreationTime   string        `json:"creationTime"`
	StartTime      string        `json:"startTime"`
	EndTime        string        `json:"endTime"`
	StoppedTime    string        `json:"stoppedTime"`
	UserID         uuid.UUID     `json:"userId"`
	UserName       string        `json:"userName"`
	UserAvatarURL  string        `json:"userAvatarURL"`
}

type CreditNoteMetaData struct {
	UserID string `json:"userId"`
}

type FindUsageResult struct {
	UsageEntries []Usage
}

// TableName sets the insert table name for this struct type
func (u *Usage) TableName() string {
	return "d_b_usage"
}

func InsertUsage(ctx context.Context, conn *gorm.DB, records ...Usage) error {
	return conn.WithContext(ctx).
		Clauses(clause.OnConflict{DoNothing: true}).
		CreateInBatches(records, 100).Error
}

func UpdateUsage(ctx context.Context, conn *gorm.DB, records ...Usage) error {
	for _, record := range records {
		err := conn.WithContext(ctx).Save(record).Error
		if err != nil {
			return fmt.Errorf("failed to update usage record ID: %s: %w", record.ID, err)
		}
	}

	return nil
}

func FindAllDraftUsage(ctx context.Context, conn *gorm.DB) ([]Usage, error) {
	var usageRecords []Usage
	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
}

type FindUsageParams struct {
	AttributionId AttributionID
	UserID        uuid.UUID
	From, To      time.Time
	ExcludeDrafts bool
	Order         Order
	Offset, Limit int64
}

func FindUsage(ctx context.Context, conn *gorm.DB, params *FindUsageParams) ([]Usage, error) {
	var usageRecords []Usage
	var usageRecordsBatch []Usage

	db := conn.WithContext(ctx).
		Where("attributionId = ?", params.AttributionId)
	if params.UserID != uuid.Nil {
		db = db.Where("metadata->>'$.userId' = ?", params.UserID.String())
	}
	db = db.Where("effectiveTime >= ? AND effectiveTime < ?", TimeToISO8601(params.From), TimeToISO8601(params.To)).
		Where("kind = ?", WorkspaceInstanceUsageKind)
	if params.ExcludeDrafts {
		db = db.Where("draft = ?", false)
	}
	db = db.Order(fmt.Sprintf("effectiveTime %s", params.Order.ToSQL()))
	if params.Offset != 0 {
		db = db.Offset(int(params.Offset))
	}
	if params.Limit != 0 {
		db = db.Limit(int(params.Limit))
	}

	result := db.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
}

type GetUsageSummaryParams struct {
	AttributionId AttributionID
	UserID        uuid.UUID
	From, To      time.Time
	ExcludeDrafts bool
}

type GetUsageSummaryResponse struct {
	CreditCentsUsed CreditCents
	NumberOfRecords int
}

func GetUsageSummary(ctx context.Context, conn *gorm.DB, params GetUsageSummaryParams) (GetUsageSummaryResponse, error) {
	db := conn.WithContext(ctx)
	query1 := db.Table((&Usage{}).TableName()).
		Select("sum(creditCents) as CreditCentsUsed, count(*) as NumberOfRecords").
		Where("attributionId = ?", params.AttributionId)
	if params.UserID != uuid.Nil {
		query1 = query1.Where("metadata->>'$.userId' = ?", params.UserID.String())
	}
	query1 = query1.Where("effectiveTime >= ? AND effectiveTime < ?", TimeToISO8601(params.From), TimeToISO8601(params.To)).
		Where("kind = ?", WorkspaceInstanceUsageKind)
	if params.ExcludeDrafts {
		query1 = query1.Where("draft = ?", false)
	}
	var result GetUsageSummaryResponse
	err := query1.Find(&result).Error
	if err != nil {
		return result, fmt.Errorf("failed to get usage meta data: %w", err)
	}
	return result, nil
}

type Balance struct {
	AttributionID AttributionID `gorm:"column:attributionId;type:varchar;size:255;" json:"attributionId"`
	CreditCents   CreditCents   `gorm:"column:creditCents;type:bigint;" json:"creditCents"`
}

func GetBalance(ctx context.Context, conn *gorm.DB, attributionId AttributionID) (CreditCents, error) {
	rows, err := conn.WithContext(ctx).
		Model(&Usage{}).
		Select("sum(creditCents) as balance").
		Where("attributionId = ?", string(attributionId)).
		Group("attributionId").
		Rows()
	if err != nil {
		return 0, fmt.Errorf("failed to get rows for list balance query: %w", err)
	}
	defer rows.Close()

	if !rows.Next() {
		return 0, nil
	}

	var balance CreditCents
	err = conn.ScanRows(rows, &balance)
	if err != nil {
		return 0, fmt.Errorf("failed to scan row: %w", err)
	}
	return balance, nil
}

func ListBalance(ctx context.Context, conn *gorm.DB) ([]Balance, error) {
	var balances []Balance
	rows, err := conn.WithContext(ctx).
		Model(&Usage{}).
		Select("attributionId as attributionId, sum(creditCents) as creditCents").
		Group("attributionId").
		Order("attributionId").
		Rows()
	if err != nil {
		return nil, fmt.Errorf("failed to get rows for list balance query: %w", err)
	}
	defer rows.Close()

	for rows.Next() {
		var balance Balance
		err = conn.ScanRows(rows, &balance)
		if err != nil {
			return nil, fmt.Errorf("failed to scan row into Balance struct: %w", err)
		}
		balances = append(balances, balance)
	}

	return balances, nil
}