diff --git a/CODEOWNERS b/CODEOWNERS index e571a160b12b7..9d97d502df6b4 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -29,3 +29,8 @@ site/src/api/countriesGenerated.ts site/src/api/rbacresourcesGenerated.ts site/src/api/typesGenerated.ts site/CLAUDE.md + +# The blood and guts of the autostop algorithm, which is quite complex and +# requires elite ball knowledge of most of the scheduling code to make changes +# without inadvertently affecting other parts of the codebase. +coderd/schedule/autostop.go @deansheather @DanielleMaywood diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go new file mode 100644 index 0000000000000..f9d54705a7cf5 --- /dev/null +++ b/coderd/database/check_constraint.go @@ -0,0 +1,16 @@ +// Code generated by scripts/dbgen/main.go. DO NOT EDIT. +package database + +// CheckConstraint represents a named check constraint on a table. +type CheckConstraint string + +// CheckConstraint enums. +const ( + CheckOneTimePasscodeSet CheckConstraint = "one_time_passcode_set" // users + CheckMaxProvisionerLogsLength CheckConstraint = "max_provisioner_logs_length" // provisioner_jobs + CheckValidationMonotonicOrder CheckConstraint = "validation_monotonic_order" // template_version_parameters + CheckMaxLogsLength CheckConstraint = "max_logs_length" // workspace_agents + CheckSubsystemsNotNone CheckConstraint = "subsystems_not_none" // workspace_agents + CheckWorkspaceBuildsAiTaskSidebarAppIDRequired CheckConstraint = "workspace_builds_ai_task_sidebar_app_id_required" // workspace_builds + CheckWorkspaceBuildsDeadlineBelowMaxDeadline CheckConstraint = "workspace_builds_deadline_below_max_deadline" // workspace_builds +) diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 053b5302d3e38..5245920ba04a9 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -2224,7 +2224,8 @@ CREATE TABLE workspace_builds ( template_version_preset_id uuid, has_ai_task boolean, ai_task_sidebar_app_id uuid, - CONSTRAINT workspace_builds_ai_task_sidebar_app_id_required CHECK (((((has_ai_task IS NULL) OR (has_ai_task = false)) AND (ai_task_sidebar_app_id IS NULL)) OR ((has_ai_task = true) AND (ai_task_sidebar_app_id IS NOT NULL)))) + CONSTRAINT workspace_builds_ai_task_sidebar_app_id_required CHECK (((((has_ai_task IS NULL) OR (has_ai_task = false)) AND (ai_task_sidebar_app_id IS NULL)) OR ((has_ai_task = true) AND (ai_task_sidebar_app_id IS NOT NULL)))), + CONSTRAINT workspace_builds_deadline_below_max_deadline CHECK ((((deadline <> '0001-01-01 00:00:00+00'::timestamp with time zone) AND (deadline <= max_deadline)) OR (max_deadline = '0001-01-01 00:00:00+00'::timestamp with time zone))) ); CREATE VIEW workspace_build_with_user AS diff --git a/coderd/database/errors.go b/coderd/database/errors.go index 0388ea2cbff49..9d0c3fee7e865 100644 --- a/coderd/database/errors.go +++ b/coderd/database/errors.go @@ -59,6 +59,28 @@ func IsForeignKeyViolation(err error, foreignKeyConstraints ...ForeignKeyConstra return false } +// IsCheckViolation checks if the error is due to a check violation. If one or +// more specific check constraints are given as arguments, the error must be +// caused by one of them. If no constraints are given, this function returns +// true for any check violation. +func IsCheckViolation(err error, checkConstraints ...CheckConstraint) bool { + var pqErr *pq.Error + if errors.As(err, &pqErr) { + if pqErr.Code.Name() == "check_violation" { + if len(checkConstraints) == 0 { + return true + } + for _, cc := range checkConstraints { + if pqErr.Constraint == string(cc) { + return true + } + } + } + } + + return false +} + // IsQueryCanceledError checks if the error is due to a query being canceled. func IsQueryCanceledError(err error) bool { var pqErr *pq.Error diff --git a/coderd/database/migrations/000356_enforce_deadline_below_max_deadline.down.sql b/coderd/database/migrations/000356_enforce_deadline_below_max_deadline.down.sql new file mode 100644 index 0000000000000..a9b2b6ff7f459 --- /dev/null +++ b/coderd/database/migrations/000356_enforce_deadline_below_max_deadline.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE workspace_builds + DROP CONSTRAINT workspace_builds_deadline_below_max_deadline; diff --git a/coderd/database/migrations/000356_enforce_deadline_below_max_deadline.up.sql b/coderd/database/migrations/000356_enforce_deadline_below_max_deadline.up.sql new file mode 100644 index 0000000000000..bcb3ab643521f --- /dev/null +++ b/coderd/database/migrations/000356_enforce_deadline_below_max_deadline.up.sql @@ -0,0 +1,20 @@ +-- New constraint: (deadline IS NOT zero AND deadline <= max_deadline) UNLESS max_deadline is zero. +-- Unfortunately, "zero" here means `time.Time{}`... + +-- Update previous builds that would fail this new constraint. This matches the +-- intended behaviour of the autostop algorithm. +UPDATE + workspace_builds +SET + deadline = max_deadline +WHERE + deadline > max_deadline + AND max_deadline != '0001-01-01 00:00:00+00'; + +-- Add the new constraint. +ALTER TABLE workspace_builds + ADD CONSTRAINT workspace_builds_deadline_below_max_deadline + CHECK ( + (deadline != '0001-01-01 00:00:00+00'::timestamptz AND deadline <= max_deadline) + OR max_deadline = '0001-01-01 00:00:00+00'::timestamptz + ); diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 9c88b9b3db679..d90967b95a384 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -6003,3 +6003,102 @@ func TestGetRunningPrebuiltWorkspaces(t *testing.T) { require.Len(t, runningPrebuilds, 1, "expected only one running prebuilt workspace") require.Equal(t, runningPrebuild.ID, runningPrebuilds[0].ID, "expected the running prebuilt workspace to be returned") } + +func TestWorkspaceBuildDeadlineConstraint(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + template := dbgen.Template(t, db, database.Template{ + CreatedBy: user.ID, + OrganizationID: org.ID, + }) + templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + TemplateID: uuid.NullUUID{UUID: template.ID, Valid: true}, + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + workspace := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + TemplateID: template.ID, + Name: "test-workspace", + Deleted: false, + }) + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + InitiatorID: database.PrebuildsSystemUserID, + Provisioner: database.ProvisionerTypeEcho, + Type: database.ProvisionerJobTypeWorkspaceBuild, + StartedAt: sql.NullTime{Time: time.Now().Add(-time.Minute), Valid: true}, + CompletedAt: sql.NullTime{Time: time.Now(), Valid: true}, + }) + workspaceBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspace.ID, + TemplateVersionID: templateVersion.ID, + JobID: job.ID, + BuildNumber: 1, + }) + + cases := []struct { + name string + deadline time.Time + maxDeadline time.Time + expectOK bool + }{ + { + name: "no deadline or max_deadline", + deadline: time.Time{}, + maxDeadline: time.Time{}, + expectOK: true, + }, + { + name: "deadline set when max_deadline is not set", + deadline: time.Now().Add(time.Hour), + maxDeadline: time.Time{}, + expectOK: true, + }, + { + name: "deadline before max_deadline", + deadline: time.Now().Add(-time.Hour), + maxDeadline: time.Now().Add(time.Hour), + expectOK: true, + }, + { + name: "deadline is max_deadline", + deadline: time.Now().Add(time.Hour), + maxDeadline: time.Now().Add(time.Hour), + expectOK: true, + }, + + { + name: "deadline after max_deadline", + deadline: time.Now().Add(time.Hour), + maxDeadline: time.Now().Add(-time.Hour), + expectOK: false, + }, + { + name: "deadline is not set when max_deadline is set", + deadline: time.Time{}, + maxDeadline: time.Now().Add(time.Hour), + expectOK: false, + }, + } + + for _, c := range cases { + err := db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ + ID: workspaceBuild.ID, + Deadline: c.deadline, + MaxDeadline: c.maxDeadline, + UpdatedAt: time.Now(), + }) + if c.expectOK { + require.NoError(t, err) + } else { + require.Error(t, err) + require.True(t, database.IsCheckViolation(err, database.CheckWorkspaceBuildsDeadlineBelowMaxDeadline)) + } + } +} diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go index 94173703c467d..d1b03cbd68a27 100644 --- a/coderd/provisionerdserver/provisionerdserver.go +++ b/coderd/provisionerdserver/provisionerdserver.go @@ -1866,8 +1866,9 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro Database: db, TemplateScheduleStore: templateScheduleStore, UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), - Now: now, - Workspace: workspace.WorkspaceTable(), + // `now` is used below to set the build completion time. + WorkspaceBuildCompletedAt: now, + Workspace: workspace.WorkspaceTable(), // Allowed to be the empty string. WorkspaceAutostart: workspace.AutostartSchedule.String, }) diff --git a/coderd/schedule/autostop.go b/coderd/schedule/autostop.go index f6a01633f3179..25bd043c60975 100644 --- a/coderd/schedule/autostop.go +++ b/coderd/schedule/autostop.go @@ -50,8 +50,19 @@ type CalculateAutostopParams struct { // by autobuild.NextAutostart WorkspaceAutostart string - Now time.Time - Workspace database.WorkspaceTable + // WorkspaceBuildCompletedAt is the time when the workspace build was + // completed. + // + // We always want to calculate using the build completion time, and not just + // the current time, to avoid forcing a workspace build's max_deadline being + // pushed to the next potential cron instance. + // + // E.g. if this function is called for an existing workspace build, which + // currently has a max_deadline within the next 2 hours (see leeway + // above), and the current time is passed into this function, the + // max_deadline will be updated to be much later than expected. + WorkspaceBuildCompletedAt time.Time + Workspace database.WorkspaceTable } type AutostopTime struct { @@ -68,8 +79,8 @@ type AutostopTime struct { // Deadline is the time when the workspace will be stopped, as long as it // doesn't see any new activity (such as SSH, app requests, etc.). When activity // is detected the deadline is bumped by the workspace's TTL (this only happens -// when activity is detected and more than 20% of the TTL has passed to save -// database queries). +// when activity is detected and more than 5% of the TTL has passed to save +// database queries, see the ActivityBumpWorkspace query). // // MaxDeadline is the maximum value for deadline. The deadline cannot be bumped // past this value, so it denotes the absolute deadline that the workspace build @@ -77,55 +88,45 @@ type AutostopTime struct { // requirement" settings and the user's "quiet hours" settings to pick a time // outside of working hours. // -// Deadline is a cost saving measure, while max deadline is a -// compliance/updating measure. +// Note that the deadline is checked at the database level: +// +// (deadline IS NOT zero AND deadline <= max_deadline) UNLESS max_deadline is zero. +// +// Deadline is intended as a cost saving measure, not as a hard policy. It is +// derived from either the workspace's TTL or the template's TTL, depending on +// the template's policy, to ensure workspaces are stopped when they are idle. +// +// MaxDeadline is intended as a compliance policy. It is derived from the +// template's autostop requirement to cap workspace uptime and effectively force +// people to update often. +// +// Note that only the build's CURRENT deadline property influences automation in +// the autobuild package. As stated above, the MaxDeadline property is only used +// to cap the value of a build's deadline. func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (AutostopTime, error) { ctx, span := tracing.StartSpan(ctx, trace.WithAttributes(attribute.String("coder.workspace_id", params.Workspace.ID.String())), trace.WithAttributes(attribute.String("coder.template_id", params.Workspace.TemplateID.String())), ) defer span.End() - defer span.End() var ( - db = params.Database - workspace = params.Workspace - now = params.Now + db = params.Database + workspace = params.Workspace + buildCompletedAt = params.WorkspaceBuildCompletedAt autostop AutostopTime ) - var ttl time.Duration - if workspace.Ttl.Valid { - // When the workspace is made it copies the template's TTL, and the user - // can unset it to disable it (unless the template has - // UserAutoStopEnabled set to false, see below). - ttl = time.Duration(workspace.Ttl.Int64) - } - - if workspace.Ttl.Valid { - // When the workspace is made it copies the template's TTL, and the user - // can unset it to disable it (unless the template has - // UserAutoStopEnabled set to false, see below). - autostop.Deadline = now.Add(time.Duration(workspace.Ttl.Int64)) - } - templateSchedule, err := params.TemplateScheduleStore.Get(ctx, db, workspace.TemplateID) if err != nil { return autostop, xerrors.Errorf("get template schedule options: %w", err) } - if !templateSchedule.UserAutostopEnabled { - // The user is not permitted to set their own TTL, so use the template - // default. - ttl = 0 - if templateSchedule.DefaultTTL > 0 { - ttl = templateSchedule.DefaultTTL - } - } + ttl := workspaceTTL(workspace, templateSchedule) if ttl > 0 { // Only apply non-zero TTLs. - autostop.Deadline = now.Add(ttl) + autostop.Deadline = buildCompletedAt.Add(ttl) if params.WorkspaceAutostart != "" { // If the deadline passes the next autostart, we need to extend the deadline to // autostart + deadline. ActivityBumpWorkspace already covers this case @@ -137,14 +138,14 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut // 3. User starts workspace at 9:45pm. // - The initial deadline is calculated to be 9:45am // - This crosses the autostart deadline, so the deadline is extended to 9pm - nextAutostart, ok := NextAutostart(params.Now, params.WorkspaceAutostart, templateSchedule) + nextAutostart, ok := NextAutostart(params.WorkspaceBuildCompletedAt, params.WorkspaceAutostart, templateSchedule) if ok && autostop.Deadline.After(nextAutostart) { autostop.Deadline = nextAutostart.Add(ttl) } } } - // Otherwise, use the autostop_requirement algorithm. + // Enforce the template autostop requirement if it's configured correctly. if templateSchedule.AutostopRequirement.DaysOfWeek != 0 { // The template has a autostop requirement, so determine the max deadline // of this workspace build. @@ -161,10 +162,10 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut // workspace. if userQuietHoursSchedule.Schedule != nil { loc := userQuietHoursSchedule.Schedule.Location() - now := now.In(loc) + buildCompletedAtInLoc := buildCompletedAt.In(loc) // Add the leeway here so we avoid checking today's quiet hours if // the workspace was started <1h before midnight. - startOfStopDay := truncateMidnight(now.Add(autostopRequirementLeeway)) + startOfStopDay := truncateMidnight(buildCompletedAtInLoc.Add(autostopRequirementLeeway)) // If the template schedule wants to only autostop on n-th weeks // then change the startOfDay to be the Monday of the next @@ -183,7 +184,7 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut // hour of the scheduled stop time will always bounce to the next // stop window). checkSchedule := userQuietHoursSchedule.Schedule.Next(startOfStopDay.Add(autostopRequirementBuffer)) - if checkSchedule.Before(now.Add(autostopRequirementLeeway)) { + if checkSchedule.Before(buildCompletedAtInLoc.Add(autostopRequirementLeeway)) { // Set the first stop day we try to tomorrow because today's // schedule is too close to now or has already passed. startOfStopDay = nextDayMidnight(startOfStopDay) @@ -213,14 +214,17 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut startOfStopDay = nextDayMidnight(startOfStopDay) } - // If the startOfDay is within an hour of now, then we add an hour. + // If the startOfDay is within an hour of the build completion time, + // then we add an hour. checkTime := startOfStopDay - if checkTime.Before(now.Add(time.Hour)) { - checkTime = now.Add(time.Hour) + if checkTime.Before(buildCompletedAtInLoc.Add(time.Hour)) { + checkTime = buildCompletedAtInLoc.Add(time.Hour) } else { - // If it's not within an hour of now, subtract 15 minutes to - // give a little leeway. This prevents skipped stop events - // because autostart perfectly lines up with autostop. + // If it's not within an hour of the build completion time, + // subtract 15 minutes to give a little leeway. This prevents + // skipped stop events because the build time (e.g. autostart + // time) perfectly lines up with the max_deadline minus the + // leeway. checkTime = checkTime.Add(autostopRequirementBuffer) } @@ -238,15 +242,35 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut autostop.Deadline = autostop.MaxDeadline } - if (!autostop.Deadline.IsZero() && autostop.Deadline.Before(now)) || (!autostop.MaxDeadline.IsZero() && autostop.MaxDeadline.Before(now)) { + if (!autostop.Deadline.IsZero() && autostop.Deadline.Before(buildCompletedAt)) || (!autostop.MaxDeadline.IsZero() && autostop.MaxDeadline.Before(buildCompletedAt)) { // Something went wrong with the deadline calculation, so we should // bail. - return autostop, xerrors.Errorf("deadline calculation error, computed deadline or max deadline is in the past for workspace build: deadline=%q maxDeadline=%q now=%q", autostop.Deadline, autostop.MaxDeadline, now) + return autostop, xerrors.Errorf("deadline calculation error, computed deadline or max deadline is in the past for workspace build: deadline=%q maxDeadline=%q now=%q", autostop.Deadline, autostop.MaxDeadline, buildCompletedAt) } return autostop, nil } +// workspaceTTL returns the TTL to use for a workspace. +// +// If the template forbids custom workspace TTLs, then we always use the +// template's configured TTL (or 0 if the template has no TTL configured). +func workspaceTTL(workspace database.WorkspaceTable, templateSchedule TemplateScheduleOptions) time.Duration { + // If the template forbids custom workspace TTLs, then we always use the + // template's configured TTL (or 0 if the template has no TTL configured). + if !templateSchedule.UserAutostopEnabled { + // This is intentionally a nested if statement because of the else if. + if templateSchedule.DefaultTTL > 0 { + return templateSchedule.DefaultTTL + } + return 0 + } + if workspace.Ttl.Valid { + return time.Duration(workspace.Ttl.Int64) + } + return 0 +} + // truncateMidnight truncates a time to midnight in the time object's timezone. // t.Truncate(24 * time.Hour) truncates based on the internal time and doesn't // factor daylight savings properly. diff --git a/coderd/schedule/autostop_test.go b/coderd/schedule/autostop_test.go index 85cc7b533a6ea..812f549f34dd2 100644 --- a/coderd/schedule/autostop_test.go +++ b/coderd/schedule/autostop_test.go @@ -76,8 +76,8 @@ func TestCalculateAutoStop(t *testing.T) { t.Log("saturdayMidnightAfterDstOut", saturdayMidnightAfterDstOut) cases := []struct { - name string - now time.Time + name string + buildCompletedAt time.Time wsAutostart string templateAutoStart schedule.TemplateAutostartRequirement @@ -98,7 +98,7 @@ func TestCalculateAutoStop(t *testing.T) { }{ { name: "OK", - now: now, + buildCompletedAt: now, templateAllowAutostop: true, templateDefaultTTL: 0, templateAutostopRequirement: schedule.TemplateAutostopRequirement{}, @@ -108,7 +108,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "Delete", - now: now, + buildCompletedAt: now, templateAllowAutostop: true, templateDefaultTTL: 0, templateAutostopRequirement: schedule.TemplateAutostopRequirement{}, @@ -118,7 +118,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "WorkspaceTTL", - now: now, + buildCompletedAt: now, templateAllowAutostop: true, templateDefaultTTL: 0, templateAutostopRequirement: schedule.TemplateAutostopRequirement{}, @@ -128,7 +128,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateDefaultTTLIgnored", - now: now, + buildCompletedAt: now, templateAllowAutostop: true, templateDefaultTTL: time.Hour, templateAutostopRequirement: schedule.TemplateAutostopRequirement{}, @@ -138,7 +138,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "WorkspaceTTLOverridesTemplateDefaultTTL", - now: now, + buildCompletedAt: now, templateAllowAutostop: true, templateDefaultTTL: 2 * time.Hour, templateAutostopRequirement: schedule.TemplateAutostopRequirement{}, @@ -148,7 +148,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateBlockWorkspaceTTL", - now: now, + buildCompletedAt: now, templateAllowAutostop: false, templateDefaultTTL: 3 * time.Hour, templateAutostopRequirement: schedule.TemplateAutostopRequirement{}, @@ -158,7 +158,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirement", - now: wednesdayMidnightUTC, + buildCompletedAt: wednesdayMidnightUTC, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -172,7 +172,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirement1HourSkip", - now: saturdayMidnightSydney.Add(-59 * time.Minute), + buildCompletedAt: saturdayMidnightSydney.Add(-59 * time.Minute), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -188,7 +188,7 @@ func TestCalculateAutoStop(t *testing.T) { // The next autostop requirement should be skipped if the // workspace is started within 1 hour of it. name: "TemplateAutostopRequirementDaily", - now: fridayEveningSydney, + buildCompletedAt: fridayEveningSydney, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -202,7 +202,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirementFortnightly/Skip", - now: wednesdayMidnightUTC, + buildCompletedAt: wednesdayMidnightUTC, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -216,7 +216,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirementFortnightly/NoSkip", - now: wednesdayMidnightUTC.AddDate(0, 0, 7), + buildCompletedAt: wednesdayMidnightUTC.AddDate(0, 0, 7), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -230,7 +230,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirementTriweekly/Skip", - now: wednesdayMidnightUTC, + buildCompletedAt: wednesdayMidnightUTC, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -246,7 +246,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirementTriweekly/NoSkip", - now: wednesdayMidnightUTC.AddDate(0, 0, 7), + buildCompletedAt: wednesdayMidnightUTC.AddDate(0, 0, 7), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -262,7 +262,7 @@ func TestCalculateAutoStop(t *testing.T) { name: "TemplateAutostopRequirementOverridesWorkspaceTTL", // now doesn't have to be UTC, but it helps us ensure that // timezones are compared correctly in this test. - now: fridayEveningSydney.In(time.UTC), + buildCompletedAt: fridayEveningSydney.In(time.UTC), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -276,7 +276,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "TemplateAutostopRequirementOverridesTemplateDefaultTTL", - now: fridayEveningSydney.In(time.UTC), + buildCompletedAt: fridayEveningSydney.In(time.UTC), templateAllowAutostop: true, templateDefaultTTL: 3 * time.Hour, userQuietHoursSchedule: sydneyQuietHours, @@ -293,7 +293,7 @@ func TestCalculateAutoStop(t *testing.T) { // The epoch is 2023-01-02 in each timezone. We set the time to // 1 second before 11pm the previous day, as this is the latest time // we allow due to our 2h leeway logic. - now: time.Date(2023, 1, 1, 21, 59, 59, 0, sydneyLoc), + buildCompletedAt: time.Date(2023, 1, 1, 21, 59, 59, 0, sydneyLoc), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -306,7 +306,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "DaylightSavings/OK", - now: duringDst, + buildCompletedAt: duringDst, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -320,7 +320,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "DaylightSavings/SwitchMidWeek/In", - now: beforeDstIn, + buildCompletedAt: beforeDstIn, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -334,7 +334,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "DaylightSavings/SwitchMidWeek/Out", - now: beforeDstOut, + buildCompletedAt: beforeDstOut, templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: sydneyQuietHours, @@ -348,7 +348,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "DaylightSavings/QuietHoursFallsOnDstSwitch/In", - now: beforeDstIn.Add(-24 * time.Hour), + buildCompletedAt: beforeDstIn.Add(-24 * time.Hour), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: dstInQuietHours, @@ -362,7 +362,7 @@ func TestCalculateAutoStop(t *testing.T) { }, { name: "DaylightSavings/QuietHoursFallsOnDstSwitch/Out", - now: beforeDstOut.Add(-24 * time.Hour), + buildCompletedAt: beforeDstOut.Add(-24 * time.Hour), templateAllowAutostop: true, templateDefaultTTL: 0, userQuietHoursSchedule: dstOutQuietHours, @@ -382,7 +382,7 @@ func TestCalculateAutoStop(t *testing.T) { // activity on the workspace. name: "AutostopCrossAutostartBorder", // Starting at 9:45pm, with the autostart at 9am. - now: pastDateNight, + buildCompletedAt: pastDateNight, templateAllowAutostop: false, templateDefaultTTL: time.Hour * 12, workspaceTTL: time.Hour * 12, @@ -405,7 +405,7 @@ func TestCalculateAutoStop(t *testing.T) { // Same as AutostopCrossAutostartBorder, but just misses the autostart. name: "AutostopCrossMissAutostartBorder", // Starting at 8:45pm, with the autostart at 9am. - now: time.Date(pastDateNight.Year(), pastDateNight.Month(), pastDateNight.Day(), 20, 30, 0, 0, chicago), + buildCompletedAt: time.Date(pastDateNight.Year(), pastDateNight.Month(), pastDateNight.Day(), 20, 30, 0, 0, chicago), templateAllowAutostop: false, templateDefaultTTL: time.Hour * 12, workspaceTTL: time.Hour * 12, @@ -429,7 +429,7 @@ func TestCalculateAutoStop(t *testing.T) { // The autostop deadline is before the autostart threshold. name: "AutostopCrossAutostartBorderMaxEarlyDeadline", // Starting at 9:45pm, with the autostart at 9am. - now: pastDateNight, + buildCompletedAt: pastDateNight, templateAllowAutostop: false, templateDefaultTTL: time.Hour * 12, workspaceTTL: time.Hour * 12, @@ -459,7 +459,7 @@ func TestCalculateAutoStop(t *testing.T) { // So the deadline is > 12 hours, but stops at the max deadline. name: "AutostopCrossAutostartBorderMaxDeadline", // Starting at 9:45pm, with the autostart at 9am. - now: pastDateNight, + buildCompletedAt: pastDateNight, templateAllowAutostop: false, templateDefaultTTL: time.Hour * 12, workspaceTTL: time.Hour * 12, @@ -571,7 +571,7 @@ func TestCalculateAutoStop(t *testing.T) { Database: db, TemplateScheduleStore: templateScheduleStore, UserQuietHoursScheduleStore: userQuietHoursScheduleStore, - Now: c.now, + WorkspaceBuildCompletedAt: c.buildCompletedAt, Workspace: workspace, WorkspaceAutostart: c.wsAutostart, }) diff --git a/coderd/workspaces.go b/coderd/workspaces.go index 2080926b44089..8d1376e7e6939 100644 --- a/coderd/workspaces.go +++ b/coderd/workspaces.go @@ -45,8 +45,8 @@ import ( ) var ( - ttlMin = time.Minute //nolint:revive // min here means 'minimum' not 'minutes' - ttlMax = 30 * 24 * time.Hour + ttlMinimum = time.Minute + ttlMaximum = 30 * 24 * time.Hour errTTLMin = xerrors.New("time until shutdown must be at least one minute") errTTLMax = xerrors.New("time until shutdown must be less than 30 days") @@ -1190,8 +1190,22 @@ func (api *API) putWorkspaceTTL(rw http.ResponseWriter, r *http.Request) { if build.Transition == database.WorkspaceTransitionStart { if err = s.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{ - ID: build.ID, - Deadline: time.Time{}, + ID: build.ID, + // Use the max_deadline as the new build deadline. It will + // either be zero (our target), or a non-zero value that we + // need to abide by anyway due to template policy. + // + // Previously, we would always set the deadline to zero, + // which was incorrect behavior. When max_deadline is + // non-zero, deadline must be set to a non-zero value that + // is less than max_deadline. + // + // Disabling TTL autostop (at a workspace or template level) + // does not trump the template's autostop requirement. + // + // Refer to the comments on schedule.CalculateAutostop for + // more information. + Deadline: build.MaxDeadline, MaxDeadline: build.MaxDeadline, UpdatedAt: dbtime.Time(api.Clock.Now()), }); err != nil { @@ -2391,11 +2405,11 @@ func validWorkspaceTTLMillis(millis *int64, templateDefault time.Duration) (sql. dur := time.Duration(*millis) * time.Millisecond truncated := dur.Truncate(time.Minute) - if truncated < ttlMin { + if truncated < ttlMinimum { return sql.NullInt64{}, errTTLMin } - if truncated > ttlMax { + if truncated > ttlMaximum { return sql.NullInt64{}, errTTLMax } diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index 9fe066aae6284..6d2a6e544ddd7 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -2897,6 +2897,55 @@ func TestWorkspaceUpdateTTL(t *testing.T) { } }) + t.Run("RemoveAutostopWithRunningWorkspaceWithMaxDeadline", func(t *testing.T) { + t.Parallel() + + var ( + ctx = testutil.Context(t, testutil.WaitLong) + client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true}) + user = coderdtest.CreateFirstUser(t, client) + version = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil) + _ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + template = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID) + deadline = 8 * time.Hour + maxDeadline = 10 * time.Hour + workspace = coderdtest.CreateWorkspace(t, client, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) { + cwr.TTLMillis = ptr.Ref(deadline.Milliseconds()) + }) + build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + ) + + // This is a hack, but the max_deadline isn't precisely configurable + // without a lot of unnecessary hassle. + dbBuild, err := db.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), build.ID) //nolint:gocritic // test + require.NoError(t, err) + dbJob, err := db.GetProvisionerJobByID(dbauthz.AsSystemRestricted(ctx), dbBuild.JobID) //nolint:gocritic // test + require.NoError(t, err) + require.True(t, dbJob.CompletedAt.Valid) + expectedMaxDeadline := dbJob.CompletedAt.Time.Add(maxDeadline) + err = db.UpdateWorkspaceBuildDeadlineByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceBuildDeadlineByIDParams{ //nolint:gocritic // test + ID: build.ID, + Deadline: dbBuild.Deadline, + MaxDeadline: expectedMaxDeadline, + UpdatedAt: dbtime.Now(), + }) + require.NoError(t, err) + + // Remove autostop. + err = client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{ + TTLMillis: nil, + }) + require.NoError(t, err) + + // Expect that the deadline is set to the max_deadline. + build, err = client.WorkspaceBuild(ctx, build.ID) + require.NoError(t, err) + require.True(t, build.Deadline.Valid) + require.WithinDuration(t, build.Deadline.Time, expectedMaxDeadline, time.Second) + require.True(t, build.MaxDeadline.Valid) + require.WithinDuration(t, build.MaxDeadline.Time, expectedMaxDeadline, time.Second) + }) + t.Run("CustomAutostopDisabledByTemplate", func(t *testing.T) { t.Parallel() var ( diff --git a/enterprise/coderd/schedule/template.go b/enterprise/coderd/schedule/template.go index 855dea4989c73..313268f2e39ad 100644 --- a/enterprise/coderd/schedule/template.go +++ b/enterprise/coderd/schedule/template.go @@ -350,14 +350,23 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Conte return nil } + // Calculate the new autostop max_deadline from the workspace. Since + // autostop is always calculated from the build completion time, we don't + // want to use the returned autostop.Deadline property as it will likely be + // in the distant past. + // + // The only exception is if the newly calculated workspace TTL is now zero, + // which means the workspace can now stay on indefinitely. + // + // This also matches the behavior of updating a workspace's TTL, where we + // don't apply the changes until the workspace is rebuilt. autostop, err := agpl.CalculateAutostop(ctx, agpl.CalculateAutostopParams{ Database: db, TemplateScheduleStore: s, UserQuietHoursScheduleStore: *s.UserQuietHoursScheduleStore.Load(), - // Use the job completion time as the time we calculate autostop from. - Now: job.CompletedAt.Time, - Workspace: workspace.WorkspaceTable(), - WorkspaceAutostart: workspace.AutostartSchedule.String, + WorkspaceBuildCompletedAt: job.CompletedAt.Time, + Workspace: workspace.WorkspaceTable(), + WorkspaceAutostart: workspace.AutostartSchedule.String, }) if err != nil { return xerrors.Errorf("calculate new autostop for workspace %q: %w", workspace.ID, err) @@ -389,9 +398,24 @@ func (s *EnterpriseTemplateScheduleStore) updateWorkspaceBuild(ctx context.Conte autostop.MaxDeadline = now.Add(time.Hour * 2) } + // If the new deadline is zero, the workspace can now stay on indefinitely. + // Otherwise, we want to discard the new value as per the comment above the + // CalculateAutostop call. + // + // We could potentially calculate a new deadline based on the TTL setting + // (on either the workspace or the template based on the template's policy) + // against the current time, but doing nothing here matches the current + // behavior of the workspace TTL update endpoint. + // + // Per the documentation of CalculateAutostop, the deadline is not intended + // as a policy measure, so it's fine that we don't update it when the + // template schedule changes. + if !autostop.Deadline.IsZero() { + autostop.Deadline = build.Deadline + } + // If the current deadline on the build is after the new max_deadline, then // set it to the max_deadline. - autostop.Deadline = build.Deadline if !autostop.MaxDeadline.IsZero() && autostop.Deadline.After(autostop.MaxDeadline) { autostop.Deadline = autostop.MaxDeadline } diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go index 4af06042b031f..feafbd54b31bb 100644 --- a/enterprise/coderd/schedule/template_test.go +++ b/enterprise/coderd/schedule/template_test.go @@ -23,7 +23,6 @@ import ( "github.com/coder/coder/v2/coderd/notifications" "github.com/coder/coder/v2/coderd/notifications/notificationstest" agplschedule "github.com/coder/coder/v2/coderd/schedule" - "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/enterprise/coderd/schedule" "github.com/coder/coder/v2/testutil" @@ -73,17 +72,23 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { buildTime := time.Date(nowY, nowM, nowD, 12, 0, 0, 0, time.UTC) // noon today UTC nextQuietHours := time.Date(nowY, nowM, nowD+1, 0, 0, 0, 0, time.UTC) // midnight tomorrow UTC - // Workspace old max_deadline too soon + defaultTTL := 8 * time.Hour + cases := []struct { - name string - now time.Time + name string + now time.Time + // Before: deadline time.Time maxDeadline time.Time - // Set to nil for no change. - newDeadline *time.Time + // After: + newDeadline time.Time newMaxDeadline time.Time - noQuietHours bool - autostopReq *agplschedule.TemplateAutostopRequirement + // Config: + noQuietHours bool + // Note that ttl will not influence the new build at all unless it's 0 + // AND the build does not have a max deadline post recalculation. + ttl time.Duration + autostopReq *agplschedule.TemplateAutostopRequirement }{ { name: "SkippedWorkspaceMaxDeadlineTooSoon", @@ -91,8 +96,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: buildTime, maxDeadline: buildTime.Add(1 * time.Hour), // Unchanged since the max deadline is too soon. - newDeadline: nil, + newDeadline: buildTime, newMaxDeadline: buildTime.Add(1 * time.Hour), + ttl: defaultTTL, // no effect }, { name: "NewWorkspaceMaxDeadlineBeforeNow", @@ -101,10 +107,11 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: buildTime, // Far into the future... maxDeadline: nextQuietHours.Add(24 * time.Hour), - newDeadline: nil, + newDeadline: buildTime, // We will use now() + 2 hours if the newly calculated max deadline // from the workspace build time is before now. newMaxDeadline: nextQuietHours.Add(8 * time.Hour), + ttl: defaultTTL, // no effect }, { name: "NewWorkspaceMaxDeadlineSoon", @@ -113,10 +120,11 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: buildTime, // Far into the future... maxDeadline: nextQuietHours.Add(24 * time.Hour), - newDeadline: nil, + newDeadline: buildTime, // We will use now() + 2 hours if the newly calculated max deadline // from the workspace build time is within the next 2 hours. newMaxDeadline: nextQuietHours.Add(1 * time.Hour), + ttl: defaultTTL, // no effect }, { name: "NewWorkspaceMaxDeadlineFuture", @@ -125,8 +133,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: buildTime, // Far into the future... maxDeadline: nextQuietHours.Add(24 * time.Hour), - newDeadline: nil, + newDeadline: buildTime, newMaxDeadline: nextQuietHours, + ttl: defaultTTL, // no effect }, { name: "DeadlineAfterNewWorkspaceMaxDeadline", @@ -136,8 +145,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: nextQuietHours.Add(24 * time.Hour), maxDeadline: nextQuietHours.Add(24 * time.Hour), // The deadline should match since it is after the new max deadline. - newDeadline: ptr.Ref(nextQuietHours), + newDeadline: nextQuietHours, newMaxDeadline: nextQuietHours, + ttl: defaultTTL, // no effect }, { // There was a bug if a user has no quiet hours set, and autostop @@ -151,13 +161,14 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: buildTime.Add(time.Hour * 8), maxDeadline: time.Time{}, // No max set // Should be unchanged - newDeadline: ptr.Ref(buildTime.Add(time.Hour * 8)), + newDeadline: buildTime.Add(time.Hour * 8), newMaxDeadline: time.Time{}, noQuietHours: true, autostopReq: &agplschedule.TemplateAutostopRequirement{ DaysOfWeek: 0, Weeks: 0, }, + ttl: defaultTTL, // no effect }, { // A bug existed where MaxDeadline could be set, but deadline was @@ -168,15 +179,15 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { deadline: time.Time{}, maxDeadline: time.Time{}, // No max set // Should be unchanged - newDeadline: ptr.Ref(time.Time{}), + newDeadline: time.Time{}, newMaxDeadline: time.Time{}, noQuietHours: true, autostopReq: &agplschedule.TemplateAutostopRequirement{ DaysOfWeek: 0, Weeks: 0, }, + ttl: defaultTTL, // no effect }, - { // Similar to 'NoDeadline' test. This has a MaxDeadline set, so // the deadline of the workspace should now be set. @@ -185,8 +196,26 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { // Start with unset times deadline: time.Time{}, maxDeadline: time.Time{}, - newDeadline: ptr.Ref(nextQuietHours), + newDeadline: nextQuietHours, newMaxDeadline: nextQuietHours, + ttl: defaultTTL, // no effect + }, + { + // If the build doesn't have a max_deadline anymore, and there is no + // TTL anymore, then both the deadline and max_deadline should be + // zero. + name: "NoTTLNoDeadlineNoMaxDeadline", + now: buildTime, + deadline: buildTime.Add(time.Hour * 8), + maxDeadline: buildTime.Add(time.Hour * 8), + newDeadline: time.Time{}, + newMaxDeadline: time.Time{}, + noQuietHours: true, + autostopReq: &agplschedule.TemplateAutostopRequirement{ + DaysOfWeek: 0, + Weeks: 0, + }, + ttl: 0, }, } @@ -206,6 +235,7 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { t.Log("maxDeadline", c.maxDeadline) t.Log("newDeadline", c.newDeadline) t.Log("newMaxDeadline", c.newMaxDeadline) + t.Log("ttl", c.ttl) var ( template = dbgen.Template(t, db, database.Template{ @@ -300,7 +330,7 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { _, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{ UserAutostartEnabled: false, UserAutostopEnabled: false, - DefaultTTL: 0, + DefaultTTL: c.ttl, AutostopRequirement: autostopReq, FailureTTL: 0, TimeTilDormant: 0, @@ -312,11 +342,8 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) { newBuild, err := db.GetWorkspaceBuildByID(ctx, wsBuild.ID) require.NoError(t, err) - if c.newDeadline == nil { - c.newDeadline = &wsBuild.Deadline - } - require.WithinDuration(t, *c.newDeadline, newBuild.Deadline, time.Second) - require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second) + require.WithinDuration(t, c.newDeadline, newBuild.Deadline, time.Second, "deadline") + require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second, "max_deadline") // Check that the new build has the same state as before. require.Equal(t, wsBuild.ProvisionerState, newBuild.ProvisionerState, "provisioner state mismatch") diff --git a/scripts/dbgen/constraint.go b/scripts/dbgen/constraint.go new file mode 100644 index 0000000000000..6853f9bb26ad5 --- /dev/null +++ b/scripts/dbgen/constraint.go @@ -0,0 +1,239 @@ +package main + +import ( + "bufio" + "bytes" + "fmt" + "os" + "path/filepath" + "regexp" + "strings" + + "golang.org/x/tools/imports" + "golang.org/x/xerrors" +) + +type constraintType string + +const ( + constraintTypeUnique constraintType = "unique" + constraintTypeForeignKey constraintType = "foreign_key" + constraintTypeCheck constraintType = "check" +) + +func (c constraintType) goType() string { + switch c { + case constraintTypeUnique: + return "UniqueConstraint" + case constraintTypeForeignKey: + return "ForeignKeyConstraint" + case constraintTypeCheck: + return "CheckConstraint" + default: + panic(fmt.Sprintf("unknown constraint type: %s", c)) + } +} + +func (c constraintType) goTypeDescriptionPart() string { + switch c { + case constraintTypeUnique: + return "unique" + case constraintTypeForeignKey: + return "foreign key" + case constraintTypeCheck: + return "check" + default: + panic(fmt.Sprintf("unknown constraint type: %s", c)) + } +} + +func (c constraintType) goEnumNamePrefix() string { + switch c { + case constraintTypeUnique: + return "Unique" + case constraintTypeForeignKey: + return "ForeignKey" + case constraintTypeCheck: + return "Check" + default: + panic(fmt.Sprintf("unknown constraint type: %s", c)) + } +} + +type constraint struct { + name string + // comment is typically the full constraint, but for check constraints it's + // instead the table name. + comment string +} + +// queryToConstraintsFn is a function that takes a query and returns zero or +// more constraints if the query matches the wanted constraint type. If the +// query does not match the wanted constraint type, the function should return +// no constraints. +type queryToConstraintsFn func(query string) ([]constraint, error) + +// generateConstraints does the following: +// 1. Read the dump.sql file +// 2. Parse the file into each query +// 3. Pass each query to the constraintFn function +// 4. Generate the enum from the returned constraints +// 5. Write the generated code to the output path +func generateConstraints(dumpPath, outputPath string, outputConstraintType constraintType, fn queryToConstraintsFn) error { + dump, err := os.Open(dumpPath) + if err != nil { + return err + } + defer dump.Close() + + var allConstraints []constraint + + dumpScanner := bufio.NewScanner(dump) + query := "" + for dumpScanner.Scan() { + line := strings.TrimSpace(dumpScanner.Text()) + switch { + case strings.HasPrefix(line, "--"): + case line == "": + case strings.HasSuffix(line, ";"): + query += line + newConstraints, err := fn(query) + query = "" + if err != nil { + return xerrors.Errorf("process query %q: %w", query, err) + } + allConstraints = append(allConstraints, newConstraints...) + default: + query += line + " " + } + } + if err = dumpScanner.Err(); err != nil { + return err + } + + s := &bytes.Buffer{} + + _, _ = fmt.Fprintf(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT. +package database + +// %[1]s represents a named %[2]s constraint on a table. +type %[1]s string + +// %[1]s enums. +const ( +`, outputConstraintType.goType(), outputConstraintType.goTypeDescriptionPart()) + + for _, c := range allConstraints { + constName := outputConstraintType.goEnumNamePrefix() + nameFromSnakeCase(c.name) + _, _ = fmt.Fprintf(s, "\t%[1]s %[2]s = %[3]q // %[4]s\n", constName, outputConstraintType.goType(), c.name, c.comment) + } + _, _ = fmt.Fprint(s, ")\n") + + data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{ + Comments: true, + }) + if err != nil { + return err + } + return os.WriteFile(outputPath, data, 0o600) +} + +// generateUniqueConstraints generates the UniqueConstraint enum. +func generateUniqueConstraints() error { + localPath, err := localFilePath() + if err != nil { + return err + } + databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database") + dumpPath := filepath.Join(databasePath, "dump.sql") + outputPath := filepath.Join(databasePath, "unique_constraint.go") + + fn := func(query string) ([]constraint, error) { + if strings.Contains(query, "UNIQUE") || strings.Contains(query, "PRIMARY KEY") { + name := "" + switch { + case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"): + name = strings.Split(query, " ")[6] + case strings.Contains(query, "CREATE UNIQUE INDEX"): + name = strings.Split(query, " ")[3] + default: + return nil, xerrors.Errorf("unknown unique constraint format: %s", query) + } + return []constraint{ + { + name: name, + comment: query, + }, + }, nil + } + return nil, nil + } + return generateConstraints(dumpPath, outputPath, constraintTypeUnique, fn) +} + +// generateForeignKeyConstraints generates the ForeignKeyConstraint enum. +func generateForeignKeyConstraints() error { + localPath, err := localFilePath() + if err != nil { + return err + } + databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database") + dumpPath := filepath.Join(databasePath, "dump.sql") + outputPath := filepath.Join(databasePath, "foreign_key_constraint.go") + + fn := func(query string) ([]constraint, error) { + if strings.Contains(query, "FOREIGN KEY") { + name := "" + switch { + case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"): + name = strings.Split(query, " ")[6] + default: + return nil, xerrors.Errorf("unknown foreign key constraint format: %s", query) + } + return []constraint{ + { + name: name, + comment: query, + }, + }, nil + } + return []constraint{}, nil + } + return generateConstraints(dumpPath, outputPath, constraintTypeForeignKey, fn) +} + +// generateCheckConstraints generates the CheckConstraint enum. +func generateCheckConstraints() error { + localPath, err := localFilePath() + if err != nil { + return err + } + databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database") + dumpPath := filepath.Join(databasePath, "dump.sql") + outputPath := filepath.Join(databasePath, "check_constraint.go") + + var ( + tableRegex = regexp.MustCompile(`CREATE TABLE\s+([^\s]+)`) + checkRegex = regexp.MustCompile(`CONSTRAINT\s+([^\s]+)\s+CHECK`) + ) + fn := func(query string) ([]constraint, error) { + constraints := []constraint{} + + tableMatches := tableRegex.FindStringSubmatch(query) + if len(tableMatches) > 0 { + table := tableMatches[1] + + // Find every CONSTRAINT xxx CHECK occurrence. + matches := checkRegex.FindAllStringSubmatch(query, -1) + for _, match := range matches { + constraints = append(constraints, constraint{ + name: match[1], + comment: table, + }) + } + } + return constraints, nil + } + + return generateConstraints(dumpPath, outputPath, constraintTypeCheck, fn) +} diff --git a/scripts/dbgen/main.go b/scripts/dbgen/main.go index 561a46199a6ef..f2f0c19b1fd0b 100644 --- a/scripts/dbgen/main.go +++ b/scripts/dbgen/main.go @@ -1,7 +1,6 @@ package main import ( - "bufio" "bytes" "fmt" "go/format" @@ -80,152 +79,12 @@ return %s return xerrors.Errorf("generate foreign key constraints: %w", err) } - return nil -} - -// generateUniqueConstraints generates the UniqueConstraint enum. -func generateUniqueConstraints() error { - localPath, err := localFilePath() - if err != nil { - return err - } - databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database") - - dump, err := os.Open(filepath.Join(databasePath, "dump.sql")) - if err != nil { - return err - } - defer dump.Close() - - var uniqueConstraints []string - dumpScanner := bufio.NewScanner(dump) - query := "" - for dumpScanner.Scan() { - line := strings.TrimSpace(dumpScanner.Text()) - switch { - case strings.HasPrefix(line, "--"): - case line == "": - case strings.HasSuffix(line, ";"): - query += line - if strings.Contains(query, "UNIQUE") || strings.Contains(query, "PRIMARY KEY") { - uniqueConstraints = append(uniqueConstraints, query) - } - query = "" - default: - query += line + " " - } - } - if err = dumpScanner.Err(); err != nil { - return err - } - - s := &bytes.Buffer{} - - _, _ = fmt.Fprint(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT. -package database -`) - _, _ = fmt.Fprint(s, ` -// UniqueConstraint represents a named unique constraint on a table. -type UniqueConstraint string - -// UniqueConstraint enums. -const ( -`) - for _, query := range uniqueConstraints { - name := "" - switch { - case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"): - name = strings.Split(query, " ")[6] - case strings.Contains(query, "CREATE UNIQUE INDEX"): - name = strings.Split(query, " ")[3] - default: - return xerrors.Errorf("unknown unique constraint format: %s", query) - } - _, _ = fmt.Fprintf(s, "\tUnique%s UniqueConstraint = %q // %s\n", nameFromSnakeCase(name), name, query) - } - _, _ = fmt.Fprint(s, ")\n") - - outputPath := filepath.Join(databasePath, "unique_constraint.go") - - data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{ - Comments: true, - }) - if err != nil { - return err - } - return os.WriteFile(outputPath, data, 0o600) -} - -// generateForeignKeyConstraints generates the ForeignKeyConstraint enum. -func generateForeignKeyConstraints() error { - localPath, err := localFilePath() - if err != nil { - return err - } - databasePath := filepath.Join(localPath, "..", "..", "..", "coderd", "database") - - dump, err := os.Open(filepath.Join(databasePath, "dump.sql")) + err = generateCheckConstraints() if err != nil { - return err - } - defer dump.Close() - - var foreignKeyConstraints []string - dumpScanner := bufio.NewScanner(dump) - query := "" - for dumpScanner.Scan() { - line := strings.TrimSpace(dumpScanner.Text()) - switch { - case strings.HasPrefix(line, "--"): - case line == "": - case strings.HasSuffix(line, ";"): - query += line - if strings.Contains(query, "FOREIGN KEY") { - foreignKeyConstraints = append(foreignKeyConstraints, query) - } - query = "" - default: - query += line + " " - } + return xerrors.Errorf("generate check constraints: %w", err) } - if err := dumpScanner.Err(); err != nil { - return err - } - - s := &bytes.Buffer{} - - _, _ = fmt.Fprint(s, `// Code generated by scripts/dbgen/main.go. DO NOT EDIT. -package database -`) - _, _ = fmt.Fprint(s, ` -// ForeignKeyConstraint represents a named foreign key constraint on a table. -type ForeignKeyConstraint string - -// ForeignKeyConstraint enums. -const ( -`) - for _, query := range foreignKeyConstraints { - name := "" - switch { - case strings.Contains(query, "ALTER TABLE") && strings.Contains(query, "ADD CONSTRAINT"): - name = strings.Split(query, " ")[6] - default: - return xerrors.Errorf("unknown foreign key constraint format: %s", query) - } - _, _ = fmt.Fprintf(s, "\tForeignKey%s ForeignKeyConstraint = %q // %s\n", nameFromSnakeCase(name), name, query) - } - _, _ = fmt.Fprint(s, ")\n") - - outputPath := filepath.Join(databasePath, "foreign_key_constraint.go") - - data, err := imports.Process(outputPath, s.Bytes(), &imports.Options{ - Comments: true, - }) - if err != nil { - return err - } - return os.WriteFile(outputPath, data, 0o600) + return nil } type stubParams struct {