opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-01 14:53:15 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
85984ff142cc937b9b15e9f314cc3385334d1b23
coder/coderd/database/dbgen/dbgen.go
T
Zach 85984ff142 feat: add enable/disable support for user secrets (#27537) 
Users can now disable a secret to stop it from being injected into
workspaces without deleting it, and re-enable it later. Disabled secrets
stay visible and editable everywhere they already appear.

An enabled secret must have at least one injection target; a secret with
no target can be stored only while disabled. Existing target-less secrets
are migrated to disabled to preserve current behavior.

Support spans the REST API, SDK, CLI, dashboard, and audit log.
2026-07-28 09:58:33 -06:00

2273 lines
96 KiB
Go
Raw Blame History

package dbgen
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/x509"
"crypto/x509/pkix"
"database/sql"
"encoding/hex"
"encoding/json"
"encoding/pem"
"errors"
"fmt"
"maps"
"math/big"
"net"
"strings"
"testing"
"time"
"github.com/google/uuid"
"github.com/sqlc-dev/pqtype"
"github.com/stretchr/testify/require"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/apikey"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/db2sdk"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/database/provisionerjobs"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/coderd/rbac/rolestore"
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/cryptorand"
"github.com/coder/coder/v2/provisionerd/proto"
"github.com/coder/coder/v2/testutil"
)
// All methods take in a 'seed' object. Any provided fields in the seed will be
// maintained. Any fields omitted will have sensible defaults generated.
// genCtx is to give all generator functions permission if the db is a dbauthz db.
var genCtx = dbauthz.As(context.Background(), rbac.Subject{
ID: "owner",
Roles: rbac.Roles(must(rbac.RoleIdentifiers{rbac.RoleOwner()}.Expand())),
Groups: []string{},
Scope: rbac.ExpandableScope(rbac.ScopeAll),
})
func AuditLog(t testing.TB, db database.Store, seed database.AuditLog) database.AuditLog {
log, err := db.InsertAuditLog(genCtx, database.InsertAuditLogParams{
ID: takeFirst(seed.ID, uuid.New()),
Time: takeFirst(seed.Time, dbtime.Now()),
UserID: takeFirst(seed.UserID, uuid.New()),
// Default to the nil uuid. So by default audit logs are not org scoped.
OrganizationID: takeFirst(seed.OrganizationID),
Ip: pqtype.Inet{
IPNet: takeFirstIP(seed.Ip.IPNet, net.IPNet{}),
Valid: takeFirst(seed.Ip.Valid, false),
},
UserAgent: sql.NullString{
String: takeFirst(seed.UserAgent.String, ""),
Valid: takeFirst(seed.UserAgent.Valid, false),
},
ResourceType: takeFirst(seed.ResourceType, database.ResourceTypeOrganization),
ResourceID: takeFirst(seed.ResourceID, uuid.New()),
ResourceTarget: takeFirst(seed.ResourceTarget, uuid.NewString()),
Action: takeFirst(seed.Action, database.AuditActionCreate),
Diff: takeFirstSlice(seed.Diff, []byte("{}")),
StatusCode: takeFirst(seed.StatusCode, 200),
AdditionalFields: takeFirstSlice(seed.AdditionalFields, []byte("{}")),
RequestID: takeFirst(seed.RequestID, uuid.New()),
ResourceIcon: takeFirst(seed.ResourceIcon, ""),
})
require.NoError(t, err, "insert audit log")
return log
}
func Chat(t testing.TB, db database.Store, seed database.Chat) database.Chat {
t.Helper()
var labels pqtype.NullRawMessage
if seed.Labels != nil {
raw, err := json.Marshal(seed.Labels)
require.NoError(t, err, "marshal chat labels")
labels = pqtype.NullRawMessage{RawMessage: raw, Valid: true}
}
chat, err := db.InsertChat(genCtx, database.InsertChatParams{
OrganizationID: takeFirst(seed.OrganizationID, uuid.New()),
OwnerID: takeFirst(seed.OwnerID, uuid.New()),
WorkspaceID: seed.WorkspaceID,
BuildID: seed.BuildID,
AgentID: seed.AgentID,
ParentChatID: seed.ParentChatID,
RootChatID: seed.RootChatID,
LastModelConfigID: takeFirst(seed.LastModelConfigID, uuid.New()),
Title: takeFirst(seed.Title, testutil.GetRandomName(t)),
Mode: seed.Mode,
PlanMode: seed.PlanMode,
Status: takeFirst(seed.Status, database.ChatStatusWaiting),
MCPServerIDs: seed.MCPServerIDs,
Labels: labels,
DynamicTools: seed.DynamicTools,
ClientType: takeFirst(seed.ClientType, database.ChatClientTypeUi),
})
require.NoError(t, err, "insert chat")
return chat
}
func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) database.ChatMessage {
t.Helper()
content := "[]"
if seed.Content.Valid {
content = string(seed.Content.RawMessage)
}
role := takeFirst(seed.Role, database.ChatMessageRoleUser)
msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{
ChatID: seed.ChatID,
CreatedBy: []uuid.UUID{seed.CreatedBy.UUID},
ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID},
ReasoningEffort: []string{string(seed.ReasoningEffort.ChatReasoningEffort)},
Role: []database.ChatMessageRole{role},
Content: []string{content},
ContentVersion: []int16{takeFirst(seed.ContentVersion, chatprompt.CurrentContentVersion)},
Visibility: []database.ChatMessageVisibility{takeFirst(seed.Visibility, database.ChatMessageVisibilityBoth)},
InputTokens: []int64{seed.InputTokens.Int64},
OutputTokens: []int64{seed.OutputTokens.Int64},
TotalTokens: []int64{seed.TotalTokens.Int64},
ReasoningTokens: []int64{seed.ReasoningTokens.Int64},
CacheCreationTokens: []int64{seed.CacheCreationTokens.Int64},
CacheReadTokens: []int64{seed.CacheReadTokens.Int64},
ContextLimit: []int64{seed.ContextLimit.Int64},
Compressed: []bool{seed.Compressed},
TotalCostMicros: []int64{seed.TotalCostMicros.Int64},
RuntimeMs: []int64{seed.RuntimeMs.Int64},
})
require.NoError(t, err, "insert chat message")
require.Len(t, msgs, 1)
return msgs[0]
}
const (
// Match the default OpenAI test model's effective context settings.
defaultChatModelContextLimit int64 = 128000
defaultChatModelCompressionThreshold int32 = 70
)
func ChatModelConfig(t testing.TB, db database.Store, seed database.ChatModelConfig, munge ...func(*database.InsertChatModelConfigParams)) database.ChatModelConfig {
t.Helper()
aiProviderID := seed.AIProviderID
if !aiProviderID.Valid {
// No AIProviderID supplied: reuse or create a default openai provider.
// Tests needing a specific provider type should pass seed.AIProviderID.
providers, err := db.GetAIProviders(genCtx, database.GetAIProvidersParams{IncludeDisabled: true})
require.NoError(t, err, "get ai providers")
var provider database.AIProvider
for _, candidate := range providers {
if candidate.Type != database.AIProviderTypeOpenai {
continue
}
if provider.ID == uuid.Nil || candidate.CreatedAt.After(provider.CreatedAt) {
provider = candidate
}
}
if provider.ID == uuid.Nil {
provider = AIProvider(t, db, database.AIProvider{
Type: database.AIProviderTypeOpenai,
})
}
aiProviderID = uuid.NullUUID{UUID: provider.ID, Valid: true}
}
params := database.InsertChatModelConfigParams{
Model: takeFirst(seed.Model, "gpt-4o-mini"),
DisplayName: takeFirst(seed.DisplayName, "Test Model"),
CreatedBy: seed.CreatedBy,
UpdatedBy: seed.UpdatedBy,
Enabled: takeFirst(seed.Enabled, true),
IsDefault: seed.IsDefault,
ContextLimit: takeFirst(seed.ContextLimit, defaultChatModelContextLimit),
CompressionThreshold: takeFirst(seed.CompressionThreshold, defaultChatModelCompressionThreshold),
Options: takeFirstSlice(seed.Options, json.RawMessage(`{}`)),
AIProviderID: aiProviderID,
}
for _, fn := range munge {
fn(&params)
}
cfg, err := db.InsertChatModelConfig(genCtx, params)
require.NoError(t, err, "insert chat model config")
return cfg
}
func AIProvider(t testing.TB, db database.Store, seed database.AIProvider, munge ...func(*database.InsertAIProviderParams)) database.AIProvider {
t.Helper()
id := seed.ID
if id == uuid.Nil {
id = uuid.New()
}
provType := seed.Type
if provType == "" {
provType = database.AIProviderTypeOpenai
}
name := takeFirst(seed.Name, testutil.GetRandomNameHyphenated(t))
displayName := seed.DisplayName
if !displayName.Valid {
displayName = sql.NullString{String: name, Valid: true}
}
params := database.InsertAIProviderParams{
ID: id,
Type: provType,
Name: name,
DisplayName: displayName,
Icon: seed.Icon,
Enabled: takeFirst(seed.Enabled, true),
// Use an unsupported scheme so leaked test provider calls fail immediately without retries.
BaseUrl: takeFirst(seed.BaseUrl, "invalid://test.invalid/"),
Settings: seed.Settings,
SettingsKeyID: seed.SettingsKeyID,
}
for _, fn := range munge {
fn(&params)
}
provider, err := db.InsertAIProvider(genCtx, params)
require.NoError(t, err, "insert ai provider")
return provider
}
func AIProviderKey(t testing.TB, db database.Store, seed database.AIProviderKey, munge ...func(*database.InsertAIProviderKeyParams)) database.AIProviderKey {
t.Helper()
id := seed.ID
if id == uuid.Nil {
id = uuid.New()
}
now := dbtime.Now()
params := database.InsertAIProviderKeyParams{
ID: id,
ProviderID: seed.ProviderID,
APIKey: takeFirst(seed.APIKey, "test-key"),
ApiKeyKeyID: seed.ApiKeyKeyID,
CreatedAt: takeFirst(seed.CreatedAt, now),
UpdatedAt: takeFirst(seed.UpdatedAt, now),
}
for _, fn := range munge {
fn(&params)
}
key, err := db.InsertAIProviderKey(genCtx, params)
require.NoError(t, err, "insert ai provider key")
return key
}
// AIProviderWithOptionalKey inserts an AI provider and, when apiKey is not
// empty, inserts a provider-scoped key for it.
func AIProviderWithOptionalKey(
t testing.TB,
db database.Store,
seed database.AIProvider,
apiKey string,
munge ...func(*database.InsertAIProviderParams),
) database.AIProvider {
t.Helper()
provider := AIProvider(t, db, seed, munge...)
if apiKey != "" {
AIProviderKey(t, db, database.AIProviderKey{
ProviderID: provider.ID,
APIKey: apiKey,
})
}
return provider
}
func ChatProvider(t testing.TB, db database.Store, seed database.ChatProvider, munge ...func(*database.InsertChatProviderParams)) database.ChatProvider {
t.Helper()
params := database.InsertChatProviderParams{
Provider: takeFirst(seed.Provider, "openai"),
DisplayName: takeFirst(seed.DisplayName, seed.Provider, "openai"),
APIKey: takeFirst(seed.APIKey, "test-key"),
BaseUrl: seed.BaseUrl,
ApiKeyKeyID: seed.ApiKeyKeyID,
CreatedBy: seed.CreatedBy,
Enabled: takeFirst(seed.Enabled, true),
CentralApiKeyEnabled: takeFirst(seed.CentralApiKeyEnabled, true),
AllowUserApiKey: seed.AllowUserApiKey,
AllowCentralApiKeyFallback: seed.AllowCentralApiKeyFallback,
}
for _, fn := range munge {
fn(&params)
}
provider := AIProvider(t, db, database.AIProvider{
Type: database.AIProviderType(params.Provider),
Name: "test-" + uuid.NewString(),
DisplayName: sql.NullString{String: params.DisplayName, Valid: params.DisplayName != ""},
BaseUrl: params.BaseUrl,
}, func(p *database.InsertAIProviderParams) {
p.Enabled = params.Enabled
})
if params.APIKey != "" {
AIProviderKey(t, db, database.AIProviderKey{
ProviderID: provider.ID,
APIKey: params.APIKey,
ApiKeyKeyID: params.ApiKeyKeyID,
})
}
return database.ChatProvider{
ID: provider.ID,
Provider: params.Provider,
DisplayName: params.DisplayName,
APIKey: params.APIKey,
BaseUrl: params.BaseUrl,
ApiKeyKeyID: params.ApiKeyKeyID,
CreatedBy: params.CreatedBy,
Enabled: params.Enabled,
CentralApiKeyEnabled: params.CentralApiKeyEnabled,
AllowUserApiKey: params.AllowUserApiKey,
AllowCentralApiKeyFallback: params.AllowCentralApiKeyFallback,
CreatedAt: provider.CreatedAt,
UpdatedAt: provider.UpdatedAt,
}
}
func MCPServerConfig(t testing.TB, db database.Store, seed database.MCPServerConfig) database.MCPServerConfig {
t.Helper()
// CreatedBy and UpdatedBy are user FKs, so default fixtures create a user.
createdBy := seed.CreatedBy.UUID
if createdBy == uuid.Nil {
createdBy = User(t, db, database.User{}).ID
}
updatedBy := seed.UpdatedBy.UUID
if updatedBy == uuid.Nil {
updatedBy = createdBy
}
cfg, err := db.InsertMCPServerConfig(genCtx, database.InsertMCPServerConfigParams{
DisplayName: takeFirst(seed.DisplayName, "Test MCP Server"),
Slug: takeFirst(seed.Slug, testutil.GetRandomName(t)),
Description: seed.Description,
IconURL: seed.IconURL,
Transport: takeFirst(seed.Transport, "streamable_http"),
Url: takeFirst(seed.Url, "https://mcp.example.com"),
AuthType: takeFirst(seed.AuthType, "none"),
OAuth2ClientID: seed.OAuth2ClientID,
OAuth2ClientSecret: seed.OAuth2ClientSecret,
OAuth2ClientSecretKeyID: seed.OAuth2ClientSecretKeyID,
OAuth2AuthURL: seed.OAuth2AuthURL,
OAuth2TokenURL: seed.OAuth2TokenURL,
OAuth2RevocationURL: seed.OAuth2RevocationURL,
OAuth2Scopes: seed.OAuth2Scopes,
APIKeyHeader: seed.APIKeyHeader,
APIKeyValue: seed.APIKeyValue,
APIKeyValueKeyID: seed.APIKeyValueKeyID,
CustomHeaders: seed.CustomHeaders,
CustomHeadersKeyID: seed.CustomHeadersKeyID,
ToolAllowList: takeFirstSlice(seed.ToolAllowList, []string{}),
ToolDenyList: takeFirstSlice(seed.ToolDenyList, []string{}),
Availability: takeFirst(seed.Availability, "default_off"),
Enabled: takeFirst(seed.Enabled, true),
ModelIntent: seed.ModelIntent,
AllowInPlanMode: seed.AllowInPlanMode,
ForwardCoderHeaders: seed.ForwardCoderHeaders,
CreatedBy: createdBy,
UpdatedBy: updatedBy,
})
require.NoError(t, err, "insert MCP server config")
return cfg
}
func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnectionLogParams) database.ConnectionLog {
arg := database.UpsertConnectionLogParams{
ID: takeFirst(seed.ID, uuid.New()),
Time: takeFirst(seed.Time, dbtime.Now()),
OrganizationID: takeFirst(seed.OrganizationID, uuid.New()),
WorkspaceOwnerID: takeFirst(seed.WorkspaceOwnerID, uuid.New()),
WorkspaceID: takeFirst(seed.WorkspaceID, uuid.New()),
WorkspaceName: takeFirst(seed.WorkspaceName, testutil.GetRandomName(t)),
AgentName: takeFirst(seed.AgentName, testutil.GetRandomName(t)),
Type: takeFirst(seed.Type, database.ConnectionTypeSsh),
Code: sql.NullInt32{
Int32: takeFirst(seed.Code.Int32, 0),
Valid: takeFirst(seed.Code.Valid, false),
},
IP: pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
},
UserAgent: sql.NullString{
String: takeFirst(seed.UserAgent.String, ""),
Valid: takeFirst(seed.UserAgent.Valid, false),
},
UserID: uuid.NullUUID{
UUID: takeFirst(seed.UserID.UUID, uuid.Nil),
Valid: takeFirst(seed.UserID.Valid, false),
},
SlugOrPort: sql.NullString{
String: takeFirst(seed.SlugOrPort.String, ""),
Valid: takeFirst(seed.SlugOrPort.Valid, false),
},
ConnectionID: uuid.NullUUID{
UUID: takeFirst(seed.ConnectionID.UUID, uuid.Nil),
Valid: takeFirst(seed.ConnectionID.Valid, false),
},
DisconnectReason: sql.NullString{
String: takeFirst(seed.DisconnectReason.String, ""),
Valid: takeFirst(seed.DisconnectReason.Valid, false),
},
ConnectionStatus: takeFirst(seed.ConnectionStatus, database.ConnectionStatusConnected),
}
var disconnectTime sql.NullTime
if arg.ConnectionStatus == database.ConnectionStatusDisconnected {
disconnectTime = sql.NullTime{Time: arg.Time, Valid: true}
}
err := db.BatchUpsertConnectionLogs(genCtx, database.BatchUpsertConnectionLogsParams{
ID: []uuid.UUID{arg.ID},
ConnectTime: []time.Time{arg.Time},
OrganizationID: []uuid.UUID{arg.OrganizationID},
WorkspaceOwnerID: []uuid.UUID{arg.WorkspaceOwnerID},
WorkspaceID: []uuid.UUID{arg.WorkspaceID},
WorkspaceName: []string{arg.WorkspaceName},
AgentName: []string{arg.AgentName},
Type: []database.ConnectionType{arg.Type},
Code: []int32{arg.Code.Int32},
CodeValid: []bool{arg.Code.Valid},
Ip: []pqtype.Inet{arg.IP},
UserAgent: []string{arg.UserAgent.String},
UserID: []uuid.UUID{arg.UserID.UUID},
SlugOrPort: []string{arg.SlugOrPort.String},
ConnectionID: []uuid.UUID{arg.ConnectionID.UUID},
DisconnectReason: []string{arg.DisconnectReason.String},
DisconnectTime: []time.Time{disconnectTime.Time},
})
require.NoError(t, err, "insert connection log")
// Query back the actual row from the database. On upsert
// conflict the DB keeps the original row's ID, so we can't
// rely on arg.ID. Match on the conflict key for rows with a
// connection_id, or by primary key for NULL connection_id.
rows, err := db.GetConnectionLogsOffset(genCtx, database.GetConnectionLogsOffsetParams{})
require.NoError(t, err, "query connection logs")
for _, row := range rows {
if arg.ConnectionID.Valid {
if row.ConnectionLog.ConnectionID == arg.ConnectionID &&
row.ConnectionLog.WorkspaceID == arg.WorkspaceID &&
row.ConnectionLog.AgentName == arg.AgentName {
return row.ConnectionLog
}
} else if row.ConnectionLog.ID == arg.ID {
return row.ConnectionLog
}
}
require.Failf(t, "connection log not found", "id=%s", arg.ID)
return database.ConnectionLog{} // unreachable
}
func BoundarySession(t testing.TB, db database.Store, seed database.BoundarySession) database.BoundarySession {
session, err := db.InsertBoundarySession(genCtx, database.InsertBoundarySessionParams{
ID: takeFirst(seed.ID, uuid.New()),
WorkspaceAgentID: takeFirst(seed.WorkspaceAgentID, uuid.New()),
OwnerID: seed.OwnerID,
ConfinedProcessName: takeFirst(seed.ConfinedProcessName, "claude-code"),
StartedAt: takeFirst(seed.StartedAt, dbtime.Now()),
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert boundary session")
return session
}
func BoundaryLogs(t testing.TB, db database.Store, seed []database.BoundaryLog) []database.BoundaryLog {
ids := make([]uuid.UUID, 0, len(seed))
sessionID := seed[0].SessionID
ownerID := seed[0].OwnerID.UUID
sequenceNumbers := make([]int32, 0, len(seed))
capturedAt := make([]time.Time, 0, len(seed))
createdAt := make([]time.Time, 0, len(seed))
protos := make([]string, 0, len(seed))
method := make([]string, 0, len(seed))
detail := make([]string, 0, len(seed))
matchedRule := make([]string, 0, len(seed))
for _, log := range seed {
log = takeFirstBoundaryLog(log)
ids = append(ids, log.ID)
sequenceNumbers = append(sequenceNumbers, log.SequenceNumber)
capturedAt = append(capturedAt, log.CapturedAt)
createdAt = append(createdAt, log.CreatedAt)
protos = append(protos, log.Proto)
method = append(method, log.Method)
detail = append(detail, log.Detail)
matchedRule = append(matchedRule, log.MatchedRule.String)
}
logs, err := db.InsertBoundaryLogs(genCtx, database.InsertBoundaryLogsParams{
ID: ids,
SessionID: sessionID,
OwnerID: ownerID,
SequenceNumber: sequenceNumbers,
CapturedAt: capturedAt,
CreatedAt: createdAt,
Proto: protos,
Method: method,
Detail: detail,
MatchedRule: matchedRule,
})
require.NoError(t, err, "insert boundary logs")
return logs
}
func takeFirstBoundaryLog(seed database.BoundaryLog) database.BoundaryLog {
seed.ID = takeFirst(seed.ID, uuid.New())
seed.SessionID = takeFirst(seed.SessionID, uuid.New())
seed.SequenceNumber = takeFirst(seed.SequenceNumber, 0)
seed.CapturedAt = takeFirst(seed.CapturedAt, dbtime.Now())
seed.CreatedAt = takeFirst(seed.CreatedAt, dbtime.Now())
seed.Proto = takeFirst(seed.Proto, "http")
seed.Method = takeFirst(seed.Method, "GET")
seed.Detail = takeFirst(seed.Detail, "https://example.com")
return seed
}
func Template(t testing.TB, db database.Store, seed database.Template) database.Template {
id := takeFirst(seed.ID, uuid.New())
if seed.GroupACL == nil {
// By default, all users in the organization can read the template.
seed.GroupACL = database.TemplateACL{
seed.OrganizationID.String(): db2sdk.TemplateRoleActions(codersdk.TemplateRoleUse),
}
}
if seed.UserACL == nil {
seed.UserACL = database.TemplateACL{}
}
err := db.InsertTemplate(genCtx, database.InsertTemplateParams{
ID: id,
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
OrganizationID: takeFirst(seed.OrganizationID, uuid.New()),
Name: takeFirst(seed.Name, testutil.GetRandomName(t)),
Provisioner: takeFirst(seed.Provisioner, database.ProvisionerTypeEcho),
ActiveVersionID: takeFirst(seed.ActiveVersionID, uuid.New()),
Description: takeFirst(seed.Description, testutil.GetRandomName(t)),
CreatedBy: takeFirst(seed.CreatedBy, uuid.New()),
Icon: takeFirst(seed.Icon, testutil.GetRandomName(t)),
UserACL: seed.UserACL,
GroupACL: seed.GroupACL,
DisplayName: takeFirst(seed.DisplayName, testutil.GetRandomName(t)),
AllowUserCancelWorkspaceJobs: seed.AllowUserCancelWorkspaceJobs,
MaxPortSharingLevel: takeFirst(seed.MaxPortSharingLevel, database.AppSharingLevelOwner),
UseClassicParameterFlow: takeFirst(seed.UseClassicParameterFlow, false),
CorsBehavior: takeFirst(seed.CorsBehavior, database.CorsBehaviorSimple),
})
require.NoError(t, err, "insert template")
template, err := db.GetTemplateByID(genCtx, id)
require.NoError(t, err, "get template")
return template
}
func APIKey(t testing.TB, db database.Store, seed database.APIKey, munge ...func(*database.InsertAPIKeyParams)) (key database.APIKey, token string) {
id, _ := cryptorand.String(10)
secret, hashed, err := apikey.GenerateSecret(22)
require.NoError(t, err)
ip := seed.IPAddress
if !ip.Valid {
ip = pqtype.Inet{
IPNet: net.IPNet{
IP: net.IPv4(127, 0, 0, 1),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
Valid: true,
}
}
// It does not make sense for the created_at to be after the expires_at.
// So if expires is set, change the default created_at to be 24 hours before.
var createdAt time.Time
if !seed.ExpiresAt.IsZero() && seed.CreatedAt.IsZero() {
createdAt = seed.ExpiresAt.Add(-24 * time.Hour)
}
params := database.InsertAPIKeyParams{
ID: takeFirst(seed.ID, id),
// 0 defaults to 86400 at the db layer
LifetimeSeconds: takeFirst(seed.LifetimeSeconds, 0),
HashedSecret: takeFirstSlice(seed.HashedSecret, hashed),
IPAddress: ip,
UserID: takeFirst(seed.UserID, uuid.New()),
LastUsed: takeFirst(seed.LastUsed, dbtime.Now()),
ExpiresAt: takeFirst(seed.ExpiresAt, dbtime.Now().Add(time.Hour)),
CreatedAt: takeFirst(seed.CreatedAt, createdAt, dbtime.Now()),
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
LoginType: takeFirst(seed.LoginType, database.LoginTypePassword),
Scopes: takeFirstSlice([]database.APIKeyScope(seed.Scopes), []database.APIKeyScope{database.ApiKeyScopeCoderAll}),
AllowList: takeFirstSlice(seed.AllowList, database.AllowList{{Type: policy.WildcardSymbol, ID: policy.WildcardSymbol}}),
TokenName: takeFirst(seed.TokenName),
}
for _, fn := range munge {
fn(&params)
}
key, err = db.InsertAPIKey(genCtx, params)
require.NoError(t, err, "insert api key")
return key, fmt.Sprintf("%s-%s", key.ID, secret)
}
func WorkspaceAgentPortShare(t testing.TB, db database.Store, orig database.WorkspaceAgentPortShare) database.WorkspaceAgentPortShare {
ps, err := db.UpsertWorkspaceAgentPortShare(genCtx, database.UpsertWorkspaceAgentPortShareParams{
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
AgentName: takeFirst(orig.AgentName, testutil.GetRandomName(t)),
Port: takeFirst(orig.Port, 8080),
ShareLevel: takeFirst(orig.ShareLevel, database.AppSharingLevelPublic),
Protocol: takeFirst(orig.Protocol, database.PortShareProtocolHttp),
})
require.NoError(t, err, "insert workspace agent")
return ps
}
func WorkspaceAgent(t testing.TB, db database.Store, orig database.WorkspaceAgent) database.WorkspaceAgent {
agt, err := db.InsertWorkspaceAgent(genCtx, database.InsertWorkspaceAgentParams{
ID: takeFirst(orig.ID, uuid.New()),
ParentID: takeFirst(orig.ParentID, uuid.NullUUID{}),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
ResourceID: takeFirst(orig.ResourceID, uuid.New()),
AuthToken: takeFirst(orig.AuthToken, uuid.New()),
AuthInstanceID: sql.NullString{
String: takeFirst(orig.AuthInstanceID.String, testutil.GetRandomName(t)),
Valid: takeFirst(orig.AuthInstanceID.Valid, true),
},
Architecture: takeFirst(orig.Architecture, "amd64"),
EnvironmentVariables: pqtype.NullRawMessage{
RawMessage: takeFirstSlice(orig.EnvironmentVariables.RawMessage, []byte("{}")),
Valid: takeFirst(orig.EnvironmentVariables.Valid, false),
},
OperatingSystem: takeFirst(orig.OperatingSystem, "linux"),
Directory: takeFirst(orig.Directory, ""),
InstanceMetadata: pqtype.NullRawMessage{
RawMessage: takeFirstSlice(orig.ResourceMetadata.RawMessage, []byte("{}")),
Valid: takeFirst(orig.ResourceMetadata.Valid, false),
},
ResourceMetadata: pqtype.NullRawMessage{
RawMessage: takeFirstSlice(orig.ResourceMetadata.RawMessage, []byte("{}")),
Valid: takeFirst(orig.ResourceMetadata.Valid, false),
},
ConnectionTimeoutSeconds: takeFirst(orig.ConnectionTimeoutSeconds, 3600),
TroubleshootingURL: takeFirst(orig.TroubleshootingURL, "https://example.com"),
MOTDFile: takeFirst(orig.MOTDFile, ""),
DisplayApps: append([]database.DisplayApp{}, orig.DisplayApps...),
DisplayOrder: takeFirst(orig.DisplayOrder, 1),
APIKeyScope: takeFirst(orig.APIKeyScope, database.AgentKeyScopeEnumAll),
})
require.NoError(t, err, "insert workspace agent")
if orig.FirstConnectedAt.Valid || orig.LastConnectedAt.Valid || orig.DisconnectedAt.Valid || orig.LastConnectedReplicaID.Valid {
err = db.UpdateWorkspaceAgentConnectionByID(genCtx, database.UpdateWorkspaceAgentConnectionByIDParams{
ID: agt.ID,
FirstConnectedAt: takeFirst(orig.FirstConnectedAt, agt.FirstConnectedAt),
LastConnectedAt: takeFirst(orig.LastConnectedAt, agt.LastConnectedAt),
DisconnectedAt: takeFirst(orig.DisconnectedAt, agt.DisconnectedAt),
LastConnectedReplicaID: takeFirst(orig.LastConnectedReplicaID, agt.LastConnectedReplicaID),
UpdatedAt: takeFirst(orig.UpdatedAt, agt.UpdatedAt),
})
require.NoError(t, err, "update workspace agent first connected at")
}
// If the lifecycle state is "ready", update the agent with the corresponding timestamps
if orig.LifecycleState == database.WorkspaceAgentLifecycleStateReady && orig.StartedAt.Valid && orig.ReadyAt.Valid {
err := db.UpdateWorkspaceAgentLifecycleStateByID(genCtx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
ID: agt.ID,
LifecycleState: orig.LifecycleState,
StartedAt: orig.StartedAt,
ReadyAt: orig.ReadyAt,
})
require.NoError(t, err, "update workspace agent lifecycle state")
}
if orig.ParentID.UUID == uuid.Nil {
// Add a test antagonist. For every agent we add a deleted sub agent
// to discover cases where deletion should be handled.
// See also `(dbfake.WorkspaceBuildBuilder).Do()`.
subAgt, err := db.InsertWorkspaceAgent(genCtx, database.InsertWorkspaceAgentParams{
ID: uuid.New(),
ParentID: uuid.NullUUID{UUID: agt.ID, Valid: true},
CreatedAt: dbtime.Now(),
UpdatedAt: dbtime.Now(),
Name: testutil.GetRandomName(t),
ResourceID: agt.ResourceID,
AuthToken: uuid.New(),
AuthInstanceID: sql.NullString{},
Architecture: agt.Architecture,
EnvironmentVariables: pqtype.NullRawMessage{},
OperatingSystem: agt.OperatingSystem,
Directory: agt.Directory,
InstanceMetadata: pqtype.NullRawMessage{},
ResourceMetadata: pqtype.NullRawMessage{},
ConnectionTimeoutSeconds: agt.ConnectionTimeoutSeconds,
TroubleshootingURL: "I AM A TEST ANTAGONIST AND I AM HERE TO MESS UP YOUR TESTS. IF YOU SEE ME, SOMETHING IS WRONG AND SUB AGENT DELETION MAY NOT BE HANDLED CORRECTLY IN A QUERY.",
MOTDFile: "",
DisplayApps: nil,
DisplayOrder: agt.DisplayOrder,
APIKeyScope: agt.APIKeyScope,
})
require.NoError(t, err, "insert workspace agent subagent antagonist")
err = db.DeleteWorkspaceSubAgentByID(genCtx, subAgt.ID)
require.NoError(t, err, "delete workspace agent subagent antagonist")
t.Logf("inserted deleted subagent antagonist %s (%v) for workspace agent %s (%v)", subAgt.Name, subAgt.ID, agt.Name, agt.ID)
}
return agt
}
func WorkspaceSubAgent(t testing.TB, db database.Store, parentAgent database.WorkspaceAgent, orig database.WorkspaceAgent) database.WorkspaceAgent {
orig.ParentID = uuid.NullUUID{UUID: parentAgent.ID, Valid: true}
orig.ResourceID = parentAgent.ResourceID
subAgt := WorkspaceAgent(t, db, orig)
return subAgt
}
func WorkspaceAgentScript(t testing.TB, db database.Store, orig database.WorkspaceAgentScript) database.WorkspaceAgentScript {
scripts, err := db.InsertWorkspaceAgentScripts(genCtx, database.InsertWorkspaceAgentScriptsParams{
WorkspaceAgentID: takeFirst(orig.WorkspaceAgentID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
LogSourceID: []uuid.UUID{takeFirst(orig.LogSourceID, uuid.New())},
LogPath: []string{takeFirst(orig.LogPath, "")},
Script: []string{takeFirst(orig.Script, "")},
Cron: []string{takeFirst(orig.Cron, "")},
StartBlocksLogin: []bool{takeFirst(orig.StartBlocksLogin, false)},
RunOnStart: []bool{takeFirst(orig.RunOnStart, false)},
RunOnStop: []bool{takeFirst(orig.RunOnStop, false)},
TimeoutSeconds: []int32{takeFirst(orig.TimeoutSeconds, 0)},
DisplayName: []string{takeFirst(orig.DisplayName, "")},
ID: []uuid.UUID{takeFirst(orig.ID, uuid.New())},
})
require.NoError(t, err, "insert workspace agent script")
require.NotEmpty(t, scripts, "insert workspace agent script returned no scripts")
return scripts[0]
}
func WorkspaceAgentScripts(t testing.TB, db database.Store, count int, orig database.WorkspaceAgentScript) []database.WorkspaceAgentScript {
scripts := make([]database.WorkspaceAgentScript, 0, count)
for range count {
scripts = append(scripts, WorkspaceAgentScript(t, db, orig))
}
return scripts
}
func WorkspaceAgentScriptTimings(t testing.TB, db database.Store, scripts []database.WorkspaceAgentScript) []database.WorkspaceAgentScriptTiming {
timings := make([]database.WorkspaceAgentScriptTiming, len(scripts))
for i, script := range scripts {
timings[i] = WorkspaceAgentScriptTiming(t, db, database.WorkspaceAgentScriptTiming{
ScriptID: script.ID,
})
}
return timings
}
func WorkspaceAgentScriptTiming(t testing.TB, db database.Store, orig database.WorkspaceAgentScriptTiming) database.WorkspaceAgentScriptTiming {
// retry a few times in case of a unique constraint violation
for i := 0; i < 10; i++ {
timing, err := db.InsertWorkspaceAgentScriptTimings(genCtx, database.InsertWorkspaceAgentScriptTimingsParams{
StartedAt: takeFirst(orig.StartedAt, dbtime.Now()),
EndedAt: takeFirst(orig.EndedAt, dbtime.Now()),
Stage: takeFirst(orig.Stage, database.WorkspaceAgentScriptTimingStageStart),
ScriptID: takeFirst(orig.ScriptID, uuid.New()),
ExitCode: takeFirst(orig.ExitCode, 0),
Status: takeFirst(orig.Status, database.WorkspaceAgentScriptTimingStatusOk),
})
if err == nil {
return timing
}
// Some tests run WorkspaceAgentScriptTiming in a loop and run into
// a unique violation - 2 rows get the same started_at value.
if (database.IsUniqueViolation(err, database.UniqueWorkspaceAgentScriptTimingsScriptIDStartedAtKey) && orig.StartedAt == time.Time{}) {
// Wait 1 millisecond so dbtime.Now() changes
time.Sleep(time.Millisecond * 1)
continue
}
require.NoError(t, err, "insert workspace agent script")
}
panic("failed to insert workspace agent script timing")
}
func WorkspaceAgentDevcontainer(t testing.TB, db database.Store, orig database.WorkspaceAgentDevcontainer) database.WorkspaceAgentDevcontainer {
devcontainers, err := db.InsertWorkspaceAgentDevcontainers(genCtx, database.InsertWorkspaceAgentDevcontainersParams{
WorkspaceAgentID: takeFirst(orig.WorkspaceAgentID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
ID: []uuid.UUID{takeFirst(orig.ID, uuid.New())},
Name: []string{takeFirst(orig.Name, testutil.GetRandomName(t))},
WorkspaceFolder: []string{takeFirst(orig.WorkspaceFolder, "/workspace")},
ConfigPath: []string{takeFirst(orig.ConfigPath, "")},
SubagentID: []uuid.UUID{orig.SubagentID.UUID},
})
require.NoError(t, err, "insert workspace agent devcontainer")
return devcontainers[0]
}
func Workspace(t testing.TB, db database.Store, orig database.WorkspaceTable) database.WorkspaceTable {
t.Helper()
var defOrgID uuid.UUID
if orig.OrganizationID == uuid.Nil {
defOrg, _ := db.GetDefaultOrganization(genCtx)
defOrgID = defOrg.ID
}
workspace, err := db.InsertWorkspace(genCtx, database.InsertWorkspaceParams{
ID: takeFirst(orig.ID, uuid.New()),
OwnerID: takeFirst(orig.OwnerID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
OrganizationID: takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
TemplateID: takeFirst(orig.TemplateID, uuid.New()),
LastUsedAt: takeFirst(orig.LastUsedAt, dbtime.Now()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
AutostartSchedule: orig.AutostartSchedule,
Ttl: orig.Ttl,
AutomaticUpdates: takeFirst(orig.AutomaticUpdates, database.AutomaticUpdatesNever),
NextStartAt: orig.NextStartAt,
})
require.NoError(t, err, "insert workspace")
if orig.Deleted {
err = db.UpdateWorkspaceDeletedByID(genCtx, database.UpdateWorkspaceDeletedByIDParams{
ID: workspace.ID,
Deleted: true,
})
require.NoError(t, err, "set workspace as deleted")
workspace.Deleted = true
}
if orig.DormantAt.Valid {
_, err = db.UpdateWorkspaceDormantDeletingAt(genCtx, database.UpdateWorkspaceDormantDeletingAtParams{
ID: workspace.ID,
DormantAt: orig.DormantAt,
})
require.NoError(t, err, "set workspace as dormant")
workspace.DormantAt = orig.DormantAt
}
if len(orig.UserACL) > 0 || len(orig.GroupACL) > 0 {
userACL := orig.UserACL
if userACL == nil {
userACL = database.WorkspaceACL{}
}
groupACL := orig.GroupACL
if groupACL == nil {
groupACL = database.WorkspaceACL{}
}
err = db.UpdateWorkspaceACLByID(genCtx, database.UpdateWorkspaceACLByIDParams{
ID: workspace.ID,
UserACL: userACL,
GroupACL: groupACL,
})
require.NoError(t, err, "set workspace ACL")
workspace.UserACL = orig.UserACL
workspace.GroupACL = orig.GroupACL
}
return workspace
}
func WorkspaceAgentLogSource(t testing.TB, db database.Store, orig database.WorkspaceAgentLogSource) database.WorkspaceAgentLogSource {
sources, err := db.InsertWorkspaceAgentLogSources(genCtx, database.InsertWorkspaceAgentLogSourcesParams{
WorkspaceAgentID: takeFirst(orig.WorkspaceAgentID, uuid.New()),
ID: []uuid.UUID{takeFirst(orig.ID, uuid.New())},
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
DisplayName: []string{takeFirst(orig.DisplayName, testutil.GetRandomName(t))},
Icon: []string{takeFirst(orig.Icon, testutil.GetRandomName(t))},
})
require.NoError(t, err, "insert workspace agent log source")
return sources[0]
}
func WorkspaceAgentLog(t testing.TB, db database.Store, orig database.WorkspaceAgentLog) database.WorkspaceAgentLog {
log, err := db.InsertWorkspaceAgentLogs(genCtx, database.InsertWorkspaceAgentLogsParams{
AgentID: takeFirst(orig.AgentID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
LogSourceID: takeFirst(orig.LogSourceID, uuid.New()),
OutputLength: int32(len(orig.Output)), // nolint: gosec // integer overflow is not a concern here
Level: []database.LogLevel{takeFirst(orig.Level, database.LogLevelInfo)},
Output: []string{takeFirst(orig.Output, "Test agent log")},
})
require.NoError(t, err, "insert workspace agent log")
require.Len(t, log, 1, "incorrect number of agent logs returned")
return log[0]
}
func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuild) database.WorkspaceBuild {
t.Helper()
buildID := takeFirst(orig.ID, uuid.New())
jobID := takeFirst(orig.JobID, uuid.New())
hasAITask := takeFirst(orig.HasAITask, sql.NullBool{})
hasExternalAgent := takeFirst(orig.HasExternalAgent, sql.NullBool{})
var build database.WorkspaceBuild
err := db.InTx(func(db database.Store) error {
err := db.InsertWorkspaceBuild(genCtx, database.InsertWorkspaceBuildParams{
ID: buildID,
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
TemplateVersionID: takeFirst(orig.TemplateVersionID, uuid.New()),
BuildNumber: takeFirst(orig.BuildNumber, 1),
Transition: takeFirst(orig.Transition, database.WorkspaceTransitionStart),
InitiatorID: takeFirst(orig.InitiatorID, uuid.New()),
JobID: jobID,
ProvisionerState: []byte{},
Deadline: takeFirst(orig.Deadline, dbtime.Now().Add(time.Hour)),
MaxDeadline: takeFirst(orig.MaxDeadline, time.Time{}),
Reason: takeFirst(orig.Reason, database.BuildReasonInitiator),
TemplateVersionPresetID: takeFirst(orig.TemplateVersionPresetID, uuid.NullUUID{
UUID: uuid.UUID{},
Valid: false,
}),
})
if err != nil {
return err
}
if orig.DailyCost > 0 {
err = db.UpdateWorkspaceBuildCostByID(genCtx, database.UpdateWorkspaceBuildCostByIDParams{
ID: buildID,
DailyCost: orig.DailyCost,
})
require.NoError(t, err)
}
if hasAITask.Valid || hasExternalAgent.Valid {
require.NoError(t, db.UpdateWorkspaceBuildFlagsByID(genCtx, database.UpdateWorkspaceBuildFlagsByIDParams{
ID: buildID,
HasAITask: hasAITask,
HasExternalAgent: hasExternalAgent,
UpdatedAt: dbtime.Now(),
}))
}
build, err = db.GetWorkspaceBuildByID(genCtx, buildID)
if err != nil {
return err
}
return nil
}, nil)
require.NoError(t, err, "insert workspace build")
return build
}
func WorkspaceBuildParameters(t testing.TB, db database.Store, orig []database.WorkspaceBuildParameter) []database.WorkspaceBuildParameter {
if len(orig) == 0 {
return nil
}
var (
names = make([]string, 0, len(orig))
values = make([]string, 0, len(orig))
params []database.WorkspaceBuildParameter
)
for _, param := range orig {
names = append(names, param.Name)
values = append(values, param.Value)
}
err := db.InTx(func(tx database.Store) error {
id := takeFirst(orig[0].WorkspaceBuildID, uuid.New())
err := tx.InsertWorkspaceBuildParameters(genCtx, database.InsertWorkspaceBuildParametersParams{
WorkspaceBuildID: id,
Name: names,
Value: values,
})
if err != nil {
return err
}
params, err = tx.GetWorkspaceBuildParameters(genCtx, id)
return err
}, nil)
require.NoError(t, err)
return params
}
func User(t testing.TB, db database.Store, orig database.User) database.User {
loginType := takeFirst(orig.LoginType, database.LoginTypePassword)
email := takeFirst(orig.Email, testutil.GetRandomName(t))
// A DB constraint requires login_type = 'none' and email = '' for service
// accounts.
if orig.IsServiceAccount {
loginType = database.LoginTypeNone
email = ""
}
user, err := db.InsertUser(genCtx, database.InsertUserParams{
ID: takeFirst(orig.ID, uuid.New()),
Email: email,
Username: takeFirst(orig.Username, testutil.GetRandomName(t)),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
HashedPassword: takeFirstSlice(orig.HashedPassword, []byte(must(cryptorand.String(32)))),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
RBACRoles: takeFirstSlice(orig.RBACRoles, []string{}),
LoginType: loginType,
Status: string(takeFirst(orig.Status, database.UserStatusDormant)),
IsServiceAccount: orig.IsServiceAccount,
})
require.NoError(t, err, "insert user")
user, err = db.UpdateUserStatus(genCtx, database.UpdateUserStatusParams{
ID: user.ID,
Status: takeFirst(orig.Status, database.UserStatusActive),
UpdatedAt: dbtime.Now(),
UserIsSeen: false,
})
require.NoError(t, err, "insert user")
if !orig.LastSeenAt.IsZero() {
user, err = db.UpdateUserLastSeenAt(genCtx, database.UpdateUserLastSeenAtParams{
ID: user.ID,
LastSeenAt: orig.LastSeenAt,
UpdatedAt: user.UpdatedAt,
})
require.NoError(t, err, "user last seen")
}
if orig.Deleted {
err = db.UpdateUserDeletedByID(genCtx, user.ID)
require.NoError(t, err, "set user as deleted")
}
return user
}
func GitSSHKey(t testing.TB, db database.Store, orig database.GitSSHKey) database.GitSSHKey {
key, err := db.InsertGitSSHKey(genCtx, database.InsertGitSSHKeyParams{
UserID: takeFirst(orig.UserID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
PrivateKey: takeFirst(orig.PrivateKey, ""),
PrivateKeyKeyID: takeFirst(orig.PrivateKeyKeyID, sql.NullString{}),
PublicKey: takeFirst(orig.PublicKey, ""),
})
require.NoError(t, err, "insert ssh key")
return key
}
func Organization(t testing.TB, db database.Store, orig database.Organization) database.Organization {
org, err := db.InsertOrganization(genCtx, database.InsertOrganizationParams{
ID: takeFirst(orig.ID, uuid.New()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
DisplayName: takeFirst(orig.Name, testutil.GetRandomName(t)),
Description: takeFirst(orig.Description, testutil.GetRandomName(t)),
Icon: takeFirst(orig.Icon, ""),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
DefaultOrgMemberRoles: takeFirstSlice(orig.DefaultOrgMemberRoles, rbac.DefaultOrgMemberRoles()),
})
require.NoError(t, err, "insert organization")
// Populate the placeholder system roles (created by DB
// trigger/migration) so org members have expected permissions.
//nolint:gocritic // ReconcileSystemRole needs the system:update
// permission that `genCtx` does not have.
sysCtx := dbauthz.AsSystemRestricted(genCtx)
for roleName := range rolestore.SystemRoleNames {
role := database.CustomRole{
Name: roleName,
OrganizationID: uuid.NullUUID{UUID: org.ID, Valid: true},
}
_, _, err = rolestore.ReconcileSystemRole(sysCtx, db, role, org)
if errors.Is(err, sql.ErrNoRows) {
// The trigger that creates the placeholder role didn't run (e.g.,
// triggers were disabled in the test). Create the role manually.
err = rolestore.CreateSystemRole(sysCtx, db, org, roleName)
require.NoError(t, err, "create role "+roleName)
_, _, err = rolestore.ReconcileSystemRole(sysCtx, db, role, org)
}
require.NoError(t, err, "reconcile role "+roleName)
}
return org
}
func OrganizationMember(t testing.TB, db database.Store, orig database.OrganizationMember) database.OrganizationMember {
mem, err := db.InsertOrganizationMember(genCtx, database.InsertOrganizationMemberParams{
OrganizationID: takeFirst(orig.OrganizationID, uuid.New()),
UserID: takeFirst(orig.UserID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
Roles: takeFirstSlice(orig.Roles, []string{}),
})
require.NoError(t, err, "insert organization")
return mem
}
func NotificationInbox(t testing.TB, db database.Store, orig database.InsertInboxNotificationParams) database.InboxNotification {
notification, err := db.InsertInboxNotification(genCtx, database.InsertInboxNotificationParams{
ID: takeFirst(orig.ID, uuid.New()),
UserID: takeFirst(orig.UserID, uuid.New()),
TemplateID: takeFirst(orig.TemplateID, uuid.New()),
Targets: takeFirstSlice(orig.Targets, []uuid.UUID{}),
Title: takeFirst(orig.Title, testutil.GetRandomName(t)),
Content: takeFirst(orig.Content, testutil.GetRandomName(t)),
Icon: takeFirst(orig.Icon, ""),
Actions: orig.Actions,
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert notification")
return notification
}
func WebpushSubscription(t testing.TB, db database.Store, orig database.InsertWebpushSubscriptionParams) database.WebpushSubscription {
subscription, err := db.InsertWebpushSubscription(genCtx, database.InsertWebpushSubscriptionParams{
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UserID: takeFirst(orig.UserID, uuid.New()),
Endpoint: takeFirst(orig.Endpoint, testutil.GetRandomName(t)),
EndpointP256dhKey: takeFirst(orig.EndpointP256dhKey, testutil.GetRandomName(t)),
EndpointAuthKey: takeFirst(orig.EndpointAuthKey, testutil.GetRandomName(t)),
})
require.NoError(t, err, "insert webpush subscription")
return subscription
}
func Group(t testing.TB, db database.Store, orig database.Group) database.Group {
t.Helper()
name := takeFirst(orig.Name, testutil.GetRandomName(t))
group, err := db.InsertGroup(genCtx, database.InsertGroupParams{
ID: takeFirst(orig.ID, uuid.New()),
Name: name,
DisplayName: takeFirst(orig.DisplayName, name),
OrganizationID: takeFirst(orig.OrganizationID, uuid.New()),
AvatarURL: takeFirst(orig.AvatarURL, "https://logo.example.com"),
QuotaAllowance: takeFirst(orig.QuotaAllowance, 0),
})
require.NoError(t, err, "insert group")
return group
}
// GroupMember requires a user + group to already exist.
// Example for creating a group member for a random group + user.
//
// GroupMember(t, db, database.GroupMemberTable{
// UserID: User(t, db, database.User{}).ID,
// GroupID: Group(t, db, database.Group{
// OrganizationID: must(db.GetDefaultOrganization(genCtx)).ID,
// }).ID,
// })
func GroupMember(t testing.TB, db database.Store, member database.GroupMemberTable) database.GroupMember {
require.NotEqual(t, member.UserID, uuid.Nil, "A user id is required to use 'dbgen.GroupMember', use 'dbgen.User'.")
require.NotEqual(t, member.GroupID, uuid.Nil, "A group id is required to use 'dbgen.GroupMember', use 'dbgen.Group'.")
//nolint:gosimple
err := db.InsertGroupMember(genCtx, database.InsertGroupMemberParams{
UserID: member.UserID,
GroupID: member.GroupID,
})
require.NoError(t, err, "insert group member")
user, err := db.GetUserByID(genCtx, member.UserID)
if errors.Is(err, sql.ErrNoRows) {
require.NoErrorf(t, err, "'dbgen.GroupMember' failed as the user with id %s does not exist. A user is required to use this function, use 'dbgen.User'.", member.UserID)
}
require.NoError(t, err, "get user by id")
group, err := db.GetGroupByID(genCtx, member.GroupID)
if errors.Is(err, sql.ErrNoRows) {
require.NoErrorf(t, err, "'dbgen.GroupMember' failed as the group with id %s does not exist. A group is required to use this function, use 'dbgen.Group'.", member.GroupID)
}
require.NoError(t, err, "get group by id")
groupMember := database.GroupMember{
UserID: user.ID,
UserEmail: user.Email,
UserUsername: user.Username,
UserHashedPassword: user.HashedPassword,
UserCreatedAt: user.CreatedAt,
UserUpdatedAt: user.UpdatedAt,
UserStatus: user.Status,
UserRbacRoles: user.RBACRoles,
UserLoginType: user.LoginType,
UserAvatarUrl: user.AvatarURL,
UserDeleted: user.Deleted,
UserLastSeenAt: user.LastSeenAt,
UserQuietHoursSchedule: user.QuietHoursSchedule,
UserName: user.Name,
UserGithubComUserID: user.GithubComUserID,
OrganizationID: group.OrganizationID,
GroupName: group.Name,
GroupID: group.ID,
}
return groupMember
}
// ProvisionerDaemon creates a provisioner daemon as far as the database is concerned. It does not run a provisioner daemon.
// If no key is provided, it will create one.
func ProvisionerDaemon(t testing.TB, db database.Store, orig database.ProvisionerDaemon) database.ProvisionerDaemon {
t.Helper()
var defOrgID uuid.UUID
if orig.OrganizationID == uuid.Nil {
defOrg, _ := db.GetDefaultOrganization(genCtx)
defOrgID = defOrg.ID
}
daemon := database.UpsertProvisionerDaemonParams{
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
OrganizationID: takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
Provisioners: takeFirstSlice(orig.Provisioners, []database.ProvisionerType{database.ProvisionerTypeEcho}),
Tags: takeFirstMap(orig.Tags, database.StringMap{}),
KeyID: takeFirst(orig.KeyID, uuid.Nil),
LastSeenAt: takeFirst(orig.LastSeenAt, sql.NullTime{Time: dbtime.Now(), Valid: true}),
Version: takeFirst(orig.Version, "v0.0.0"),
APIVersion: takeFirst(orig.APIVersion, "1.1"),
}
if daemon.KeyID == uuid.Nil {
key, err := db.InsertProvisionerKey(genCtx, database.InsertProvisionerKeyParams{
ID: uuid.New(),
Name: daemon.Name + "-key",
OrganizationID: daemon.OrganizationID,
HashedSecret: []byte("secret"),
CreatedAt: dbtime.Now(),
Tags: daemon.Tags,
})
require.NoError(t, err)
daemon.KeyID = key.ID
}
d, err := db.UpsertProvisionerDaemon(genCtx, daemon)
require.NoError(t, err)
return d
}
// ProvisionerJob is a bit more involved to get the values such as "completedAt", "startedAt", "cancelledAt" set. ps
// can be set to nil if you are SURE that you don't require a provisionerdaemon to acquire the job in your test.
func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig database.ProvisionerJob) database.ProvisionerJob {
t.Helper()
var defOrgID uuid.UUID
if orig.OrganizationID == uuid.Nil {
defOrg, _ := db.GetDefaultOrganization(genCtx)
defOrgID = defOrg.ID
}
jobID := takeFirst(orig.ID, uuid.New())
// Always set some tags to prevent Acquire from grabbing jobs it should not.
tags := maps.Clone(orig.Tags)
if !orig.StartedAt.Time.IsZero() {
if tags == nil {
tags = make(database.StringMap)
}
// Make sure when we acquire the job, we only get this one.
tags[jobID.String()] = "true"
}
job, err := db.InsertProvisionerJob(genCtx, database.InsertProvisionerJobParams{
ID: jobID,
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
OrganizationID: takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
InitiatorID: takeFirst(orig.InitiatorID, uuid.New()),
Provisioner: takeFirst(orig.Provisioner, database.ProvisionerTypeEcho),
StorageMethod: takeFirst(orig.StorageMethod, database.ProvisionerStorageMethodFile),
FileID: takeFirst(orig.FileID, uuid.New()),
Type: takeFirst(orig.Type, database.ProvisionerJobTypeWorkspaceBuild),
Input: takeFirstSlice(orig.Input, []byte("{}")),
Tags: tags,
TraceMetadata: pqtype.NullRawMessage{},
LogsOverflowed: false,
})
require.NoError(t, err, "insert job")
if ps != nil {
err = provisionerjobs.PostJob(ps, job)
require.NoError(t, err, "post job to pubsub")
}
if !orig.StartedAt.Time.IsZero() {
job, err = db.AcquireProvisionerJob(genCtx, database.AcquireProvisionerJobParams{
StartedAt: orig.StartedAt,
OrganizationID: job.OrganizationID,
Types: []database.ProvisionerType{job.Provisioner},
ProvisionerTags: must(json.Marshal(tags)),
WorkerID: takeFirst(orig.WorkerID, uuid.NullUUID{}),
})
require.NoError(t, err)
// There is no easy way to make sure we acquire the correct job.
require.Equal(t, jobID, job.ID, "acquired incorrect job")
}
if !orig.CompletedAt.Time.IsZero() || orig.Error.String != "" {
err = db.UpdateProvisionerJobWithCompleteByID(genCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
ID: jobID,
UpdatedAt: job.UpdatedAt,
CompletedAt: orig.CompletedAt,
Error: orig.Error,
ErrorCode: orig.ErrorCode,
})
require.NoError(t, err)
}
if !orig.CanceledAt.Time.IsZero() {
err = db.UpdateProvisionerJobWithCancelByID(genCtx, database.UpdateProvisionerJobWithCancelByIDParams{
ID: jobID,
CanceledAt: orig.CanceledAt,
CompletedAt: orig.CompletedAt,
})
require.NoError(t, err)
}
job, err = db.GetProvisionerJobByID(genCtx, jobID)
require.NoError(t, err, "get job: %s", jobID.String())
return job
}
func ProvisionerJobLog(t testing.TB, db database.Store, orig database.ProvisionerJobLog) database.ProvisionerJobLog {
logs, err := db.InsertProvisionerJobLogs(genCtx, database.InsertProvisionerJobLogsParams{
JobID: takeFirst(orig.JobID, uuid.New()),
CreatedAt: []time.Time{takeFirst(orig.CreatedAt, dbtime.Now())},
Source: []database.LogSource{takeFirst(orig.Source, database.LogSourceProvisioner)},
Level: []database.LogLevel{takeFirst(orig.Level, database.LogLevelInfo)},
Stage: []string{takeFirst(orig.Stage, "Test")},
Output: []string{takeFirst(orig.Output, "Provisioner job log")},
})
require.NoError(t, err, "insert provisioner job log")
require.Len(t, logs, 1, "insert provisioner job log returned incorrect number of logs")
return logs[0]
}
func ProvisionerKey(t testing.TB, db database.Store, orig database.ProvisionerKey) database.ProvisionerKey {
key, err := db.InsertProvisionerKey(genCtx, database.InsertProvisionerKeyParams{
ID: takeFirst(orig.ID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
OrganizationID: takeFirst(orig.OrganizationID, uuid.New()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
HashedSecret: orig.HashedSecret,
Tags: orig.Tags,
})
require.NoError(t, err, "insert provisioner key")
return key
}
func WorkspaceApp(t testing.TB, db database.Store, orig database.WorkspaceApp) database.WorkspaceApp {
resource, err := db.UpsertWorkspaceApp(genCtx, database.UpsertWorkspaceAppParams{
ID: takeFirst(orig.ID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
AgentID: takeFirst(orig.AgentID, uuid.New()),
Slug: takeFirst(orig.Slug, testutil.GetRandomName(t)),
DisplayName: takeFirst(orig.DisplayName, testutil.GetRandomName(t)),
Icon: takeFirst(orig.Icon, testutil.GetRandomName(t)),
Command: sql.NullString{
String: takeFirst(orig.Command.String, "ls"),
Valid: orig.Command.Valid,
},
Url: sql.NullString{
String: takeFirst(orig.Url.String),
Valid: orig.Url.Valid,
},
External: orig.External,
Subdomain: orig.Subdomain,
SharingLevel: takeFirst(orig.SharingLevel, database.AppSharingLevelOwner),
HealthcheckUrl: takeFirst(orig.HealthcheckUrl, "https://localhost:8000"),
HealthcheckInterval: takeFirst(orig.HealthcheckInterval, 60),
HealthcheckThreshold: takeFirst(orig.HealthcheckThreshold, 60),
Health: takeFirst(orig.Health, database.WorkspaceAppHealthHealthy),
DisplayOrder: takeFirst(orig.DisplayOrder, 1),
DisplayGroup: orig.DisplayGroup,
Hidden: orig.Hidden,
OpenIn: takeFirst(orig.OpenIn, database.WorkspaceAppOpenInSlimWindow),
Tooltip: takeFirst(orig.Tooltip, testutil.GetRandomName(t)),
})
require.NoError(t, err, "insert app")
return resource
}
func WorkspaceAppStat(t testing.TB, db database.Store, orig database.WorkspaceAppStat) database.WorkspaceAppStat {
// This is not going to be correct, but our query doesn't return the ID.
id, err := cryptorand.Int63()
require.NoError(t, err, "generate id")
scheme := database.WorkspaceAppStat{
ID: takeFirst(orig.ID, id),
UserID: takeFirst(orig.UserID, uuid.New()),
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
AgentID: takeFirst(orig.AgentID, uuid.New()),
AccessMethod: takeFirst(orig.AccessMethod, ""),
SlugOrPort: takeFirst(orig.SlugOrPort, ""),
SessionID: takeFirst(orig.SessionID, uuid.New()),
SessionStartedAt: takeFirst(orig.SessionStartedAt, dbtime.Now().Add(-time.Minute)),
SessionEndedAt: takeFirst(orig.SessionEndedAt, dbtime.Now()),
Requests: takeFirst(orig.Requests, 1),
}
err = db.InsertWorkspaceAppStats(genCtx, database.InsertWorkspaceAppStatsParams{
UserID: []uuid.UUID{scheme.UserID},
WorkspaceID: []uuid.UUID{scheme.WorkspaceID},
AgentID: []uuid.UUID{scheme.AgentID},
AccessMethod: []string{scheme.AccessMethod},
SlugOrPort: []string{scheme.SlugOrPort},
SessionID: []uuid.UUID{scheme.SessionID},
SessionStartedAt: []time.Time{scheme.SessionStartedAt},
SessionEndedAt: []time.Time{scheme.SessionEndedAt},
Requests: []int32{scheme.Requests},
})
require.NoError(t, err, "insert workspace agent stat")
return scheme
}
func WorkspaceAppStatus(t testing.TB, db database.Store, orig database.WorkspaceAppStatus) database.WorkspaceAppStatus {
appStatus, err := db.InsertWorkspaceAppStatus(genCtx, database.InsertWorkspaceAppStatusParams{
ID: takeFirst(orig.ID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
AgentID: takeFirst(orig.AgentID, uuid.New()),
AppID: takeFirst(orig.AppID, uuid.New()),
State: takeFirst(orig.State, database.WorkspaceAppStatusStateWorking),
Message: takeFirst(orig.Message, ""),
Uri: takeFirst(orig.Uri, sql.NullString{}),
})
require.NoError(t, err, "insert workspace agent status")
return appStatus
}
func WorkspaceResource(t testing.TB, db database.Store, orig database.WorkspaceResource) database.WorkspaceResource {
resource, err := db.InsertWorkspaceResource(genCtx, database.InsertWorkspaceResourceParams{
ID: takeFirst(orig.ID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
JobID: takeFirst(orig.JobID, uuid.New()),
Transition: takeFirst(orig.Transition, database.WorkspaceTransitionStart),
Type: takeFirst(orig.Type, "fake_resource"),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
Hide: takeFirst(orig.Hide, false),
Icon: takeFirst(orig.Icon, ""),
InstanceType: sql.NullString{
String: takeFirst(orig.InstanceType.String, ""),
Valid: takeFirst(orig.InstanceType.Valid, false),
},
DailyCost: takeFirst(orig.DailyCost, 0),
ModulePath: sql.NullString{
String: takeFirst(orig.ModulePath.String, ""),
Valid: takeFirst(orig.ModulePath.Valid, true),
},
})
require.NoError(t, err, "insert resource")
return resource
}
func WorkspaceModule(t testing.TB, db database.Store, orig database.WorkspaceModule) database.WorkspaceModule {
module, err := db.InsertWorkspaceModule(genCtx, database.InsertWorkspaceModuleParams{
ID: takeFirst(orig.ID, uuid.New()),
JobID: takeFirst(orig.JobID, uuid.New()),
Transition: takeFirst(orig.Transition, database.WorkspaceTransitionStart),
Source: takeFirst(orig.Source, "test-source"),
Version: takeFirst(orig.Version, "v1.0.0"),
Key: takeFirst(orig.Key, "test-key"),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert workspace module")
return module
}
func WorkspaceResourceMetadatums(t testing.TB, db database.Store, seed database.WorkspaceResourceMetadatum) []database.WorkspaceResourceMetadatum {
meta, err := db.InsertWorkspaceResourceMetadata(genCtx, database.InsertWorkspaceResourceMetadataParams{
WorkspaceResourceID: takeFirst(seed.WorkspaceResourceID, uuid.New()),
Key: []string{takeFirst(seed.Key, testutil.GetRandomName(t))},
Value: []string{takeFirst(seed.Value.String, testutil.GetRandomName(t))},
Sensitive: []bool{takeFirst(seed.Sensitive, false)},
})
require.NoError(t, err, "insert meta data")
return meta
}
func WorkspaceProxy(t testing.TB, db database.Store, orig database.WorkspaceProxy) (database.WorkspaceProxy, string) {
secret, hashedSecret, err := apikey.GenerateSecret(64)
require.NoError(t, err, "generate secret")
proxy, err := db.InsertWorkspaceProxy(genCtx, database.InsertWorkspaceProxyParams{
ID: takeFirst(orig.ID, uuid.New()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
DisplayName: takeFirst(orig.DisplayName, testutil.GetRandomName(t)),
Icon: takeFirst(orig.Icon, testutil.GetRandomName(t)),
TokenHashedSecret: hashedSecret,
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
DerpEnabled: takeFirst(orig.DerpEnabled, false),
DerpOnly: takeFirst(orig.DerpEnabled, false),
})
require.NoError(t, err, "insert proxy")
// Also set these fields if the caller wants them.
if orig.Url != "" || orig.WildcardHostname != "" {
proxy, err = db.RegisterWorkspaceProxy(genCtx, database.RegisterWorkspaceProxyParams{
Url: orig.Url,
WildcardHostname: orig.WildcardHostname,
ID: proxy.ID,
})
require.NoError(t, err, "update proxy")
}
return proxy, secret
}
func File(t testing.TB, db database.Store, orig database.File) database.File {
file, err := db.InsertFile(genCtx, database.InsertFileParams{
ID: takeFirst(orig.ID, uuid.New()),
Hash: takeFirst(orig.Hash, hex.EncodeToString(make([]byte, 32))),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
CreatedBy: takeFirst(orig.CreatedBy, uuid.New()),
Mimetype: takeFirst(orig.Mimetype, "application/x-tar"),
Data: takeFirstSlice(orig.Data, []byte{}),
})
require.NoError(t, err, "insert file")
return file
}
func UserLink(t testing.TB, db database.Store, orig database.UserLink) database.UserLink {
link, err := db.InsertUserLink(genCtx, database.InsertUserLinkParams{
UserID: takeFirst(orig.UserID, uuid.New()),
LoginType: takeFirst(orig.LoginType, database.LoginTypeGithub),
LinkedID: takeFirst(orig.LinkedID),
OAuthAccessToken: takeFirst(orig.OAuthAccessToken, uuid.NewString()),
OAuthAccessTokenKeyID: takeFirst(orig.OAuthAccessTokenKeyID, sql.NullString{}),
OAuthRefreshToken: takeFirst(orig.OAuthRefreshToken, uuid.NewString()),
OAuthRefreshTokenKeyID: takeFirst(orig.OAuthRefreshTokenKeyID, sql.NullString{}),
OAuthExpiry: takeFirst(orig.OAuthExpiry, dbtime.Now().Add(time.Hour*24)),
Claims: orig.Claims,
})
require.NoError(t, err, "insert link")
return link
}
func ExternalAuthLink(t testing.TB, db database.Store, orig database.ExternalAuthLink) database.ExternalAuthLink {
msg := takeFirst(&orig.OAuthExtra, &pqtype.NullRawMessage{})
link, err := db.InsertExternalAuthLink(genCtx, database.InsertExternalAuthLinkParams{
ProviderID: takeFirst(orig.ProviderID, uuid.New().String()),
UserID: takeFirst(orig.UserID, uuid.New()),
OAuthAccessToken: takeFirst(orig.OAuthAccessToken, uuid.NewString()),
OAuthAccessTokenKeyID: takeFirst(orig.OAuthAccessTokenKeyID, sql.NullString{}),
OAuthRefreshToken: takeFirst(orig.OAuthRefreshToken, uuid.NewString()),
OAuthRefreshTokenKeyID: takeFirst(orig.OAuthRefreshTokenKeyID, sql.NullString{}),
OAuthExpiry: takeFirst(orig.OAuthExpiry, dbtime.Now().Add(time.Hour*24)),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
OAuthExtra: *msg,
})
require.NoError(t, err, "insert external auth link")
return link
}
func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVersion) database.TemplateVersion {
var version database.TemplateVersion
hasAITask := takeFirst(orig.HasAITask, sql.NullBool{})
hasExternalAgent := takeFirst(orig.HasExternalAgent, sql.NullBool{})
jobID := takeFirst(orig.JobID, uuid.New())
err := db.InTx(func(db database.Store) error {
versionID := takeFirst(orig.ID, uuid.New())
err := db.InsertTemplateVersion(genCtx, database.InsertTemplateVersionParams{
ID: versionID,
TemplateID: takeFirst(orig.TemplateID, uuid.NullUUID{}),
OrganizationID: takeFirst(orig.OrganizationID, uuid.New()),
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
Message: orig.Message,
Readme: takeFirst(orig.Readme, testutil.GetRandomName(t)),
JobID: jobID,
CreatedBy: takeFirst(orig.CreatedBy, uuid.New()),
SourceExampleID: takeFirst(orig.SourceExampleID, sql.NullString{}),
})
if err != nil {
return err
}
if hasAITask.Valid || hasExternalAgent.Valid {
require.NoError(t, db.UpdateTemplateVersionFlagsByJobID(genCtx, database.UpdateTemplateVersionFlagsByJobIDParams{
JobID: jobID,
HasAITask: hasAITask,
HasExternalAgent: hasExternalAgent,
UpdatedAt: dbtime.Now(),
}))
}
version, err = db.GetTemplateVersionByID(genCtx, versionID)
if err != nil {
return err
}
return nil
}, nil)
require.NoError(t, err, "insert template version")
return version
}
func TemplateVersionVariable(t testing.TB, db database.Store, orig database.TemplateVersionVariable) database.TemplateVersionVariable {
version, err := db.InsertTemplateVersionVariable(genCtx, database.InsertTemplateVersionVariableParams{
TemplateVersionID: takeFirst(orig.TemplateVersionID, uuid.New()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
Description: takeFirst(orig.Description, testutil.GetRandomName(t)),
Type: takeFirst(orig.Type, "string"),
Value: takeFirst(orig.Value, ""),
DefaultValue: takeFirst(orig.DefaultValue, testutil.GetRandomName(t)),
Required: takeFirst(orig.Required, false),
Sensitive: takeFirst(orig.Sensitive, false),
})
require.NoError(t, err, "insert template version variable")
return version
}
func TemplateVersionWorkspaceTag(t testing.TB, db database.Store, orig database.TemplateVersionWorkspaceTag) database.TemplateVersionWorkspaceTag {
workspaceTag, err := db.InsertTemplateVersionWorkspaceTag(genCtx, database.InsertTemplateVersionWorkspaceTagParams{
TemplateVersionID: takeFirst(orig.TemplateVersionID, uuid.New()),
Key: takeFirst(orig.Key, testutil.GetRandomName(t)),
Value: takeFirst(orig.Value, testutil.GetRandomName(t)),
})
require.NoError(t, err, "insert template version workspace tag")
return workspaceTag
}
func TemplateVersionParameter(t testing.TB, db database.Store, orig database.TemplateVersionParameter) database.TemplateVersionParameter {
t.Helper()
version, err := db.InsertTemplateVersionParameter(genCtx, database.InsertTemplateVersionParameterParams{
TemplateVersionID: takeFirst(orig.TemplateVersionID, uuid.New()),
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
Description: takeFirst(orig.Description, testutil.GetRandomName(t)),
Type: takeFirst(orig.Type, "string"),
FormType: orig.FormType, // empty string is ok!
Mutable: takeFirst(orig.Mutable, false),
DefaultValue: takeFirst(orig.DefaultValue, testutil.GetRandomName(t)),
Icon: takeFirst(orig.Icon, testutil.GetRandomName(t)),
Options: takeFirstSlice(orig.Options, []byte("[]")),
ValidationRegex: takeFirst(orig.ValidationRegex, ""),
ValidationMin: takeFirst(orig.ValidationMin, sql.NullInt32{}),
ValidationMax: takeFirst(orig.ValidationMax, sql.NullInt32{}),
ValidationError: takeFirst(orig.ValidationError, ""),
ValidationMonotonic: takeFirst(orig.ValidationMonotonic, ""),
Required: takeFirst(orig.Required, false),
DisplayName: takeFirst(orig.DisplayName, testutil.GetRandomName(t)),
DisplayOrder: takeFirst(orig.DisplayOrder, 0),
Ephemeral: takeFirst(orig.Ephemeral, false),
})
require.NoError(t, err, "insert template version parameter")
return version
}
func TemplateVersionTerraformValues(t testing.TB, db database.Store, orig database.TemplateVersionTerraformValue) database.TemplateVersionTerraformValue {
t.Helper()
jobID := uuid.New()
if orig.TemplateVersionID != uuid.Nil {
v, err := db.GetTemplateVersionByID(genCtx, orig.TemplateVersionID)
if err == nil {
jobID = v.JobID
}
}
params := database.InsertTemplateVersionTerraformValuesByJobIDParams{
JobID: jobID,
CachedPlan: takeFirstSlice(orig.CachedPlan, []byte("{}")),
CachedModuleFiles: orig.CachedModuleFiles,
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
ProvisionerdVersion: takeFirst(orig.ProvisionerdVersion, proto.CurrentVersion.String()),
}
err := db.InsertTemplateVersionTerraformValuesByJobID(genCtx, params)
require.NoError(t, err, "insert template version parameter")
v, err := db.GetTemplateVersionTerraformValues(genCtx, orig.TemplateVersionID)
require.NoError(t, err, "get template version values")
return v
}
func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.WorkspaceAgentStat) database.WorkspaceAgentStat {
if orig.ConnectionsByProto == nil {
orig.ConnectionsByProto = json.RawMessage([]byte("{}"))
}
jsonProto := []byte(fmt.Sprintf("[%s]", orig.ConnectionsByProto))
params := database.InsertWorkspaceAgentStatsParams{
ID: []uuid.UUID{takeFirst(orig.ID, uuid.New())},
CreatedAt: []time.Time{takeFirst(orig.CreatedAt, dbtime.Now())},
UserID: []uuid.UUID{takeFirst(orig.UserID, uuid.New())},
TemplateID: []uuid.UUID{takeFirst(orig.TemplateID, uuid.New())},
WorkspaceID: []uuid.UUID{takeFirst(orig.WorkspaceID, uuid.New())},
AgentID: []uuid.UUID{takeFirst(orig.AgentID, uuid.New())},
ConnectionsByProto: jsonProto,
ConnectionCount: []int64{takeFirst(orig.ConnectionCount, 0)},
RxPackets: []int64{takeFirst(orig.RxPackets, 0)},
RxBytes: []int64{takeFirst(orig.RxBytes, 0)},
TxPackets: []int64{takeFirst(orig.TxPackets, 0)},
TxBytes: []int64{takeFirst(orig.TxBytes, 0)},
SessionCountVSCode: []int64{takeFirst(orig.SessionCountVSCode, 0)},
SessionCountJetBrains: []int64{takeFirst(orig.SessionCountJetBrains, 0)},
SessionCountReconnectingPTY: []int64{takeFirst(orig.SessionCountReconnectingPTY, 0)},
SessionCountSSH: []int64{takeFirst(orig.SessionCountSSH, 0)},
ConnectionMedianLatencyMS: []float64{takeFirst(orig.ConnectionMedianLatencyMS, 0)},
Usage: []bool{takeFirst(orig.Usage, false)},
}
err := db.InsertWorkspaceAgentStats(genCtx, params)
require.NoError(t, err, "insert workspace agent stat")
return database.WorkspaceAgentStat{
ID: params.ID[0],
CreatedAt: params.CreatedAt[0],
UserID: params.UserID[0],
AgentID: params.AgentID[0],
WorkspaceID: params.WorkspaceID[0],
TemplateID: params.TemplateID[0],
ConnectionsByProto: orig.ConnectionsByProto,
ConnectionCount: params.ConnectionCount[0],
RxPackets: params.RxPackets[0],
RxBytes: params.RxBytes[0],
TxPackets: params.TxPackets[0],
TxBytes: params.TxBytes[0],
ConnectionMedianLatencyMS: params.ConnectionMedianLatencyMS[0],
SessionCountVSCode: params.SessionCountVSCode[0],
SessionCountJetBrains: params.SessionCountJetBrains[0],
SessionCountReconnectingPTY: params.SessionCountReconnectingPTY[0],
SessionCountSSH: params.SessionCountSSH[0],
Usage: params.Usage[0],
}
}
func OAuth2ProviderApp(t testing.TB, db database.Store, seed database.OAuth2ProviderApp) database.OAuth2ProviderApp {
app, err := db.InsertOAuth2ProviderApp(genCtx, database.InsertOAuth2ProviderAppParams{
ID: takeFirst(seed.ID, uuid.New()),
Name: takeFirst(seed.Name, testutil.GetRandomName(t)),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
Icon: takeFirst(seed.Icon, ""),
CallbackURL: takeFirst(seed.CallbackURL, "http://localhost"),
RedirectUris: takeFirstSlice(seed.RedirectUris, []string{}),
ClientType: takeFirst(seed.ClientType, sql.NullString{String: "confidential", Valid: true}),
DynamicallyRegistered: takeFirst(seed.DynamicallyRegistered, sql.NullBool{Bool: false, Valid: true}),
ClientIDIssuedAt: takeFirst(seed.ClientIDIssuedAt, sql.NullTime{}),
ClientSecretExpiresAt: takeFirst(seed.ClientSecretExpiresAt, sql.NullTime{}),
GrantTypes: takeFirstSlice(seed.GrantTypes, []string{"authorization_code", "refresh_token"}),
ResponseTypes: takeFirstSlice(seed.ResponseTypes, []string{"code"}),
TokenEndpointAuthMethod: takeFirst(seed.TokenEndpointAuthMethod, sql.NullString{String: "client_secret_basic", Valid: true}),
Scope: takeFirst(seed.Scope, sql.NullString{}),
Contacts: takeFirstSlice(seed.Contacts, []string{}),
ClientUri: takeFirst(seed.ClientUri, sql.NullString{}),
LogoUri: takeFirst(seed.LogoUri, sql.NullString{}),
TosUri: takeFirst(seed.TosUri, sql.NullString{}),
PolicyUri: takeFirst(seed.PolicyUri, sql.NullString{}),
JwksUri: takeFirst(seed.JwksUri, sql.NullString{}),
Jwks: seed.Jwks, // pqtype.NullRawMessage{} is not comparable, use existing value
SoftwareID: takeFirst(seed.SoftwareID, sql.NullString{}),
SoftwareVersion: takeFirst(seed.SoftwareVersion, sql.NullString{}),
RegistrationAccessToken: seed.RegistrationAccessToken,
RegistrationClientUri: takeFirst(seed.RegistrationClientUri, sql.NullString{}),
})
require.NoError(t, err, "insert oauth2 app")
return app
}
func OAuth2ProviderAppSecret(t testing.TB, db database.Store, seed database.OAuth2ProviderAppSecret) database.OAuth2ProviderAppSecret {
app, err := db.InsertOAuth2ProviderAppSecret(genCtx, database.InsertOAuth2ProviderAppSecretParams{
ID: takeFirst(seed.ID, uuid.New()),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
SecretPrefix: takeFirstSlice(seed.SecretPrefix, []byte("prefix")),
HashedSecret: takeFirstSlice(seed.HashedSecret, []byte("hashed-secret")),
DisplaySecret: takeFirst(seed.DisplaySecret, "secret"),
AppID: takeFirst(seed.AppID, uuid.New()),
})
require.NoError(t, err, "insert oauth2 app secret")
return app
}
func OAuth2ProviderAppCode(t testing.TB, db database.Store, seed database.OAuth2ProviderAppCode) database.OAuth2ProviderAppCode {
code, err := db.InsertOAuth2ProviderAppCode(genCtx, database.InsertOAuth2ProviderAppCodeParams{
ID: takeFirst(seed.ID, uuid.New()),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
ExpiresAt: takeFirst(seed.CreatedAt, dbtime.Now()),
SecretPrefix: takeFirstSlice(seed.SecretPrefix, []byte("prefix")),
HashedSecret: takeFirstSlice(seed.HashedSecret, []byte("hashed-secret")),
AppID: takeFirst(seed.AppID, uuid.New()),
UserID: takeFirst(seed.UserID, uuid.New()),
ResourceUri: seed.ResourceUri,
CodeChallenge: seed.CodeChallenge,
CodeChallengeMethod: seed.CodeChallengeMethod,
StateHash: seed.StateHash,
RedirectUri: seed.RedirectUri,
})
require.NoError(t, err, "insert oauth2 app code")
return code
}
func OAuth2ProviderAppToken(t testing.TB, db database.Store, seed database.OAuth2ProviderAppToken) database.OAuth2ProviderAppToken {
token, err := db.InsertOAuth2ProviderAppToken(genCtx, database.InsertOAuth2ProviderAppTokenParams{
ID: takeFirst(seed.ID, uuid.New()),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
ExpiresAt: takeFirst(seed.CreatedAt, dbtime.Now()),
HashPrefix: takeFirstSlice(seed.HashPrefix, []byte("prefix")),
RefreshHash: takeFirstSlice(seed.RefreshHash, []byte("hashed-secret")),
AppSecretID: takeFirst(seed.AppSecretID, uuid.New()),
APIKeyID: takeFirst(seed.APIKeyID, uuid.New().String()),
UserID: takeFirst(seed.UserID, uuid.New()),
Audience: seed.Audience,
})
require.NoError(t, err, "insert oauth2 app token")
return token
}
func WorkspaceAgentMemoryResourceMonitor(t testing.TB, db database.Store, seed database.WorkspaceAgentMemoryResourceMonitor) database.WorkspaceAgentMemoryResourceMonitor {
monitor, err := db.InsertMemoryResourceMonitor(genCtx, database.InsertMemoryResourceMonitorParams{
AgentID: takeFirst(seed.AgentID, uuid.New()),
Enabled: takeFirst(seed.Enabled, true),
State: takeFirst(seed.State, database.WorkspaceAgentMonitorStateOK),
Threshold: takeFirst(seed.Threshold, 100),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
DebouncedUntil: takeFirst(seed.DebouncedUntil, time.Time{}),
})
require.NoError(t, err, "insert workspace agent memory resource monitor")
return monitor
}
func WorkspaceAgentVolumeResourceMonitor(t testing.TB, db database.Store, seed database.WorkspaceAgentVolumeResourceMonitor) database.WorkspaceAgentVolumeResourceMonitor {
monitor, err := db.InsertVolumeResourceMonitor(genCtx, database.InsertVolumeResourceMonitorParams{
AgentID: takeFirst(seed.AgentID, uuid.New()),
Path: takeFirst(seed.Path, "/"),
Enabled: takeFirst(seed.Enabled, true),
State: takeFirst(seed.State, database.WorkspaceAgentMonitorStateOK),
Threshold: takeFirst(seed.Threshold, 100),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
UpdatedAt: takeFirst(seed.UpdatedAt, dbtime.Now()),
DebouncedUntil: takeFirst(seed.DebouncedUntil, time.Time{}),
})
require.NoError(t, err, "insert workspace agent volume resource monitor")
return monitor
}
func CustomRole(t testing.TB, db database.Store, seed database.CustomRole) database.CustomRole {
role, err := db.InsertCustomRole(genCtx, database.InsertCustomRoleParams{
Name: takeFirst(seed.Name, strings.ToLower(testutil.GetRandomName(t))),
DisplayName: testutil.GetRandomName(t),
OrganizationID: seed.OrganizationID,
SitePermissions: takeFirstSlice(seed.SitePermissions, []database.CustomRolePermission{}),
OrgPermissions: takeFirstSlice(seed.SitePermissions, []database.CustomRolePermission{}),
UserPermissions: takeFirstSlice(seed.SitePermissions, []database.CustomRolePermission{}),
MemberPermissions: takeFirstSlice(seed.MemberPermissions, []database.CustomRolePermission{}),
IsSystem: seed.IsSystem,
})
require.NoError(t, err, "insert custom role")
return role
}
func CryptoKey(t testing.TB, db database.Store, seed database.CryptoKey) database.CryptoKey {
t.Helper()
seed.Feature = takeFirst(seed.Feature, database.CryptoKeyFeatureWorkspaceAppsAPIKey)
// An empty string for the secret is interpreted as
// a caller wanting a new secret to be generated.
// To generate a key with a NULL secret set Valid=false
// and String to a non-empty string.
if seed.Secret.String == "" {
secret, err := newCryptoKeySecret(seed.Feature)
require.NoError(t, err, "generate secret")
seed.Secret = sql.NullString{
String: secret,
Valid: true,
}
}
key, err := db.InsertCryptoKey(genCtx, database.InsertCryptoKeyParams{
Sequence: takeFirst(seed.Sequence, 123),
Secret: seed.Secret,
SecretKeyID: takeFirst(seed.SecretKeyID, sql.NullString{}),
Feature: seed.Feature,
StartsAt: takeFirst(seed.StartsAt, dbtime.Now()),
})
require.NoError(t, err, "insert crypto key")
if seed.DeletesAt.Valid {
key, err = db.UpdateCryptoKeyDeletesAt(genCtx, database.UpdateCryptoKeyDeletesAtParams{
Feature: key.Feature,
Sequence: key.Sequence,
DeletesAt: sql.NullTime{Time: seed.DeletesAt.Time, Valid: true},
})
require.NoError(t, err, "update crypto key deletes_at")
}
return key
}
func ProvisionerJobTimings(t testing.TB, db database.Store, build database.WorkspaceBuild, count int) []database.ProvisionerJobTiming {
timings := make([]database.ProvisionerJobTiming, count)
for i := range count {
timings[i] = provisionerJobTiming(t, db, database.ProvisionerJobTiming{
JobID: build.JobID,
})
}
return timings
}
func TelemetryItem(t testing.TB, db database.Store, seed database.TelemetryItem) database.TelemetryItem {
if seed.Key == "" {
seed.Key = testutil.GetRandomName(t)
}
if seed.Value == "" {
seed.Value = time.Now().Format(time.RFC3339)
}
err := db.UpsertTelemetryItem(genCtx, database.UpsertTelemetryItemParams{
Key: seed.Key,
Value: seed.Value,
})
require.NoError(t, err, "upsert telemetry item")
item, err := db.GetTelemetryItem(genCtx, seed.Key)
require.NoError(t, err, "get telemetry item")
return item
}
func Preset(t testing.TB, db database.Store, seed database.InsertPresetParams) database.TemplateVersionPreset {
preset, err := db.InsertPreset(genCtx, database.InsertPresetParams{
ID: takeFirst(seed.ID, uuid.New()),
TemplateVersionID: takeFirst(seed.TemplateVersionID, uuid.New()),
Name: takeFirst(seed.Name, testutil.GetRandomName(t)),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
DesiredInstances: seed.DesiredInstances,
InvalidateAfterSecs: seed.InvalidateAfterSecs,
SchedulingTimezone: seed.SchedulingTimezone,
IsDefault: seed.IsDefault,
Description: seed.Description,
Icon: seed.Icon,
LastInvalidatedAt: seed.LastInvalidatedAt,
})
require.NoError(t, err, "insert preset")
return preset
}
func PresetPrebuildSchedule(t testing.TB, db database.Store, seed database.InsertPresetPrebuildScheduleParams) database.TemplateVersionPresetPrebuildSchedule {
schedule, err := db.InsertPresetPrebuildSchedule(genCtx, database.InsertPresetPrebuildScheduleParams{
PresetID: takeFirst(seed.PresetID, uuid.New()),
CronExpression: takeFirst(seed.CronExpression, "* 9-18 * * 1-5"),
DesiredInstances: takeFirst(seed.DesiredInstances, 1),
})
require.NoError(t, err, "insert preset prebuild schedule")
return schedule
}
func PresetParameter(t testing.TB, db database.Store, seed database.InsertPresetParametersParams) []database.TemplateVersionPresetParameter {
parameters, err := db.InsertPresetParameters(genCtx, database.InsertPresetParametersParams{
TemplateVersionPresetID: takeFirst(seed.TemplateVersionPresetID, uuid.New()),
Names: takeFirstSlice(seed.Names, []string{testutil.GetRandomName(t)}),
Values: takeFirstSlice(seed.Values, []string{testutil.GetRandomName(t)}),
})
require.NoError(t, err, "insert preset parameters")
return parameters
}
func UserSecret(t testing.TB, db database.Store, seed database.UserSecret, mutators ...func(params *database.CreateUserSecretParams)) database.UserSecret {
params := database.CreateUserSecretParams{
ID: takeFirst(seed.ID, uuid.New()),
UserID: takeFirst(seed.UserID, uuid.New()),
Name: takeFirst(seed.Name, "secret-name"),
Description: takeFirst(seed.Description, "secret description"),
Value: takeFirst(seed.Value, "secret value"),
ValueKeyID: seed.ValueKeyID,
EnvName: takeFirst(seed.EnvName, "SECRET_ENV_NAME"),
FilePath: takeFirst(seed.FilePath, "~/secret/file/path"),
Enabled: takeFirst(seed.Enabled, true),
}
for _, mut := range mutators {
mut(&params)
}
userSecret, err := db.CreateUserSecret(genCtx, params)
require.NoError(t, err, "failed to insert user secret")
return userSecret
}
func ClaimPrebuild(
t testing.TB,
db database.Store,
now time.Time,
newUserID uuid.UUID,
newName string,
presetID uuid.UUID,
autostartSchedule sql.NullString,
nextStartAt sql.NullTime,
ttl sql.NullInt64,
) database.ClaimPrebuiltWorkspaceRow {
claimedWorkspace, err := db.ClaimPrebuiltWorkspace(genCtx, database.ClaimPrebuiltWorkspaceParams{
NewUserID: newUserID,
NewName: newName,
Now: now,
PresetID: presetID,
AutostartSchedule: autostartSchedule,
NextStartAt: nextStartAt,
WorkspaceTtl: ttl,
})
require.NoError(t, err, "claim prebuilt workspace")
return claimedWorkspace
}
func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertAIBridgeInterceptionParams, endedAt *time.Time) database.AIBridgeInterception {
interception, err := db.InsertAIBridgeInterception(genCtx, database.InsertAIBridgeInterceptionParams{
ID: takeFirst(seed.ID, uuid.New()),
APIKeyID: seed.APIKeyID,
InitiatorID: takeFirst(seed.InitiatorID, uuid.New()),
Provider: takeFirst(seed.Provider, "provider"),
ProviderName: takeFirst(seed.ProviderName, "provider-name"),
Model: takeFirst(seed.Model, "model"),
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
StartedAt: takeFirst(seed.StartedAt, dbtime.Now()),
Client: seed.Client,
ThreadParentInterceptionID: seed.ThreadParentInterceptionID,
ThreadRootInterceptionID: seed.ThreadRootInterceptionID,
ClientSessionID: seed.ClientSessionID,
CredentialKind: takeFirst(seed.CredentialKind, database.CredentialKindCentralized),
CredentialHint: takeFirst(seed.CredentialHint, ""),
AgentFirewallSessionID: seed.AgentFirewallSessionID,
AgentFirewallSequenceNumber: seed.AgentFirewallSequenceNumber,
})
if endedAt != nil {
interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{
ID: interception.ID,
EndedAt: *endedAt,
CredentialHint: takeFirst(seed.CredentialHint, ""),
ErrorType: database.NullAIBridgeInterceptionErrorType{},
ErrorMessage: sql.NullString{},
})
require.NoError(t, err, "insert aibridge interception")
}
require.NoError(t, err, "insert aibridge interception")
return interception
}
func AIBridgeTokenUsage(t testing.TB, db database.Store, seed database.InsertAIBridgeTokenUsageParams) database.AIBridgeTokenUsage {
usage, err := db.InsertAIBridgeTokenUsage(genCtx, database.InsertAIBridgeTokenUsageParams{
ID: takeFirst(seed.ID, uuid.New()),
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
InputTokens: takeFirst(seed.InputTokens, 100),
OutputTokens: takeFirst(seed.OutputTokens, 100),
CacheReadInputTokens: seed.CacheReadInputTokens,
CacheWriteInputTokens: seed.CacheWriteInputTokens,
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
EffectiveGroupID: seed.EffectiveGroupID,
InputPriceMicros: seed.InputPriceMicros,
OutputPriceMicros: seed.OutputPriceMicros,
CacheReadPriceMicros: seed.CacheReadPriceMicros,
CacheWritePriceMicros: seed.CacheWritePriceMicros,
CostMicros: seed.CostMicros,
})
require.NoError(t, err, "insert aibridge token usage")
return usage
}
func AIBridgeUserPrompt(t testing.TB, db database.Store, seed database.InsertAIBridgeUserPromptParams) database.AIBridgeUserPrompt {
prompt, err := db.InsertAIBridgeUserPrompt(genCtx, database.InsertAIBridgeUserPromptParams{
ID: takeFirst(seed.ID, uuid.New()),
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
Prompt: takeFirst(seed.Prompt, "prompt"),
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert aibridge user prompt")
return prompt
}
func AIBridgeToolUsage(t testing.TB, db database.Store, seed database.InsertAIBridgeToolUsageParams) database.AIBridgeToolUsage {
serverURL := sql.NullString{}
if seed.ServerUrl.Valid {
serverURL = seed.ServerUrl
}
invocationError := sql.NullString{}
if seed.InvocationError.Valid {
invocationError = seed.InvocationError
}
toolUsage, err := db.InsertAIBridgeToolUsage(genCtx, database.InsertAIBridgeToolUsageParams{
ID: takeFirst(seed.ID, uuid.New()),
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
ProviderToolCallID: takeFirst(seed.ProviderToolCallID),
ProviderItemID: takeFirst(seed.ProviderItemID),
Tool: takeFirst(seed.Tool, "tool"),
ServerUrl: serverURL,
Input: takeFirst(seed.Input, "input"),
Injected: takeFirst(seed.Injected, false),
InvocationError: invocationError,
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert aibridge tool usage")
return toolUsage
}
func AIBridgeModelThought(t testing.TB, db database.Store, seed database.InsertAIBridgeModelThoughtParams) database.AIBridgeModelThought {
thought, err := db.InsertAIBridgeModelThought(genCtx, database.InsertAIBridgeModelThoughtParams{
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
Content: takeFirst(seed.Content, ""),
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert aibridge model thought")
return thought
}
func Task(t testing.TB, db database.Store, orig database.TaskTable) database.Task {
t.Helper()
parameters := orig.TemplateParameters
if parameters == nil {
parameters = json.RawMessage([]byte("{}"))
}
task, err := db.InsertTask(genCtx, database.InsertTaskParams{
ID: takeFirst(orig.ID, uuid.New()),
OrganizationID: orig.OrganizationID,
OwnerID: orig.OwnerID,
Name: takeFirst(orig.Name, testutil.GetRandomNameHyphenated(t)),
DisplayName: takeFirst(orig.DisplayName, testutil.GetRandomNameHyphenated(t)),
WorkspaceID: orig.WorkspaceID,
TemplateVersionID: orig.TemplateVersionID,
TemplateParameters: parameters,
Prompt: orig.Prompt,
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "failed to insert task")
// Return the Task from the view instead of the TaskTable
fetched, err := db.GetTaskByID(genCtx, task.ID)
require.NoError(t, err, "failed to fetch task")
require.Equal(t, task.ID, fetched.ID)
return fetched
}
func TaskWorkspaceApp(t testing.TB, db database.Store, orig database.TaskWorkspaceApp) database.TaskWorkspaceApp {
t.Helper()
app, err := db.UpsertTaskWorkspaceApp(genCtx, database.UpsertTaskWorkspaceAppParams{
TaskID: orig.TaskID,
WorkspaceBuildNumber: orig.WorkspaceBuildNumber,
WorkspaceAgentID: orig.WorkspaceAgentID,
WorkspaceAppID: orig.WorkspaceAppID,
})
require.NoError(t, err, "failed to upsert task workspace app")
return app
}
func provisionerJobTiming(t testing.TB, db database.Store, seed database.ProvisionerJobTiming) database.ProvisionerJobTiming {
timing, err := db.InsertProvisionerJobTimings(genCtx, database.InsertProvisionerJobTimingsParams{
JobID: takeFirst(seed.JobID, uuid.New()),
StartedAt: []time.Time{takeFirst(seed.StartedAt, dbtime.Now())},
EndedAt: []time.Time{takeFirst(seed.EndedAt, dbtime.Now())},
Stage: []database.ProvisionerJobTimingStage{takeFirst(seed.Stage, database.ProvisionerJobTimingStageInit)},
Source: []string{takeFirst(seed.Source, "source")},
Action: []string{takeFirst(seed.Action, "action")},
Resource: []string{takeFirst(seed.Resource, "resource")},
})
require.NoError(t, err, "insert provisioner job timing")
return timing[0]
}
func must[V any](v V, err error) V {
if err != nil {
panic(err)
}
return v
}
func takeFirstIP(values ...net.IPNet) net.IPNet {
return takeFirstF(values, func(v net.IPNet) bool {
return len(v.IP) != 0 && len(v.Mask) != 0
})
}
// takeFirstSlice implements takeFirst for []any.
// []any is not a comparable type.
func takeFirstSlice[T any](values ...[]T) []T {
return takeFirstF(values, func(v []T) bool {
return len(v) != 0
})
}
func takeFirstMap[T, E comparable](values ...map[T]E) map[T]E {
return takeFirstF(values, func(v map[T]E) bool {
return v != nil
})
}
// takeFirstF takes the first value that returns true
func takeFirstF[Value any](values []Value, take func(v Value) bool) Value {
for _, v := range values {
if take(v) {
return v
}
}
// If all empty, return the last element
if len(values) > 0 {
return values[len(values)-1]
}
var empty Value
return empty
}
// takeFirst will take the first non-empty value.
func takeFirst[Value comparable](values ...Value) Value {
var empty Value
return takeFirstF(values, func(v Value) bool {
return v != empty
})
}
func newCryptoKeySecret(feature database.CryptoKeyFeature) (string, error) {
switch feature {
case database.CryptoKeyFeatureWorkspaceAppsAPIKey:
return generateCryptoKey(32)
case database.CryptoKeyFeatureWorkspaceAppsToken:
return generateCryptoKey(64)
case database.CryptoKeyFeatureOIDCConvert:
return generateCryptoKey(64)
case database.CryptoKeyFeatureTailnetResume:
return generateCryptoKey(64)
case database.CryptoKeyFeatureNATSCA:
return generateCACryptoKeySecret()
}
return "", xerrors.Errorf("unknown feature: %s", feature)
}
// generateCACryptoKeySecret generates a self-signed CA certificate and private
// key as a PEM bundle, matching the secret format that coderd/cryptokeys
// produces for the nats_ca feature. It intentionally duplicates
// cryptokeys.generateCASecret rather than calling it: coderd/cryptokeys's
// internal tests import dbgen, so importing cryptokeys here would create a
// test-build import cycle.
func generateCACryptoKeySecret() (string, error) {
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
return "", xerrors.Errorf("generate key: %w", err)
}
template := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "dbgen-ca"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(24 * time.Hour),
KeyUsage: x509.KeyUsageCertSign,
BasicConstraintsValid: true,
IsCA: true,
MaxPathLenZero: true,
}
der, err := x509.CreateCertificate(rand.Reader, template, template, key.Public(), key)
if err != nil {
return "", xerrors.Errorf("create certificate: %w", err)
}
keyDER, err := x509.MarshalECPrivateKey(key)
if err != nil {
return "", xerrors.Errorf("marshal private key: %w", err)
}
var secret []byte
secret = append(secret, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})...)
secret = append(secret, pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})...)
return string(secret), nil
}
func generateCryptoKey(length int) (string, error) {
b := make([]byte, length)
_, err := rand.Read(b)
if err != nil {
return "", xerrors.Errorf("rand read: %w", err)
}
return hex.EncodeToString(b), nil
}
Reference in New Issue View Git Blame Copy Permalink