From 17cc8811d0f6d6da04e160c1f3836ae356b087ab Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Fri, 1 Nov 2024 11:33:55 -0400 Subject: [PATCH 1/5] chore: add unit test for testing workspace quota query --- coderd/database/dbgen/dbgen.go | 9 +++ enterprise/coderd/workspacequota_test.go | 99 +++++++++++++++++++++++- 2 files changed, 106 insertions(+), 2 deletions(-) diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 4ac675309f662..5e83125a93b84 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -288,6 +288,15 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil if err != nil { return err } + + if orig.DailyCost > 0 { + err = db.UpdateWorkspaceBuildCostByID(genCtx, database.UpdateWorkspaceBuildCostByIDParams{ + ID: buildID, + DailyCost: orig.DailyCost, + }) + require.NoError(t, err) + } + build, err = db.GetWorkspaceBuildByID(genCtx, buildID) if err != nil { return err diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 13142f11e5717..93dd7d0d84aa7 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -21,6 +21,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" @@ -31,9 +32,13 @@ import ( ) func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, organizationID string, consumed, total int) { + verifyQuotaUser(ctx, t, client, organizationID, codersdk.Me, consumed, total) +} + +func verifyQuotaUser(ctx context.Context, t *testing.T, client *codersdk.Client, organizationID string, user string, consumed, total int) { t.Helper() - got, err := client.WorkspaceQuota(ctx, organizationID, codersdk.Me) + got, err := client.WorkspaceQuota(ctx, organizationID, user) require.NoError(t, err) require.EqualValues(t, codersdk.WorkspaceQuota{ Budget: total, @@ -43,7 +48,7 @@ func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, org // Remove this check when the deprecated endpoint is removed. // This just makes sure the deprecated endpoint is still working // as intended. It will only work for the default organization. - deprecatedGot, err := deprecatedQuotaEndpoint(ctx, client, codersdk.Me) + deprecatedGot, err := deprecatedQuotaEndpoint(ctx, client, user) require.NoError(t, err, "deprecated endpoint") // Only continue to check if the values differ if deprecatedGot.Budget != got.Budget || deprecatedGot.CreditsConsumed != got.CreditsConsumed { @@ -300,6 +305,96 @@ func TestWorkspaceQuota(t *testing.T) { verifyQuota(ctx, t, owner, first.OrganizationID.String(), 0, 30) verifyQuota(ctx, t, owner, second.ID.String(), 0, 15) }) + + // ManyWorkspaces uses dbfake and dbgen to insert a scenario into the db. + t.Run("ManyWorkspaces", func(t *testing.T) { + t.Parallel() + + owner, db, first := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureTemplateRBAC: 1, + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + client, _ := coderdtest.CreateAnotherUser(t, owner, first.OrganizationID, rbac.RoleOwner()) + + // Prepopulate database. Use dbfake as it is quicker and + // easier than the api. + ctx := testutil.Context(t, testutil.WaitLong) + ctx = dbauthz.AsSystemRestricted(ctx) + + user := dbgen.User(t, db, database.User{}) + noise := dbgen.User(t, db, database.User{}) + + second := dbfake.Organization(t, db). + Members(user, noise). + EveryoneAllowance(10). + Group(database.Group{ + QuotaAllowance: 25, + }, user, noise). + Group(database.Group{ + QuotaAllowance: 30, + }, noise). + Do() + + third := dbfake.Organization(t, db). + Members(noise). + EveryoneAllowance(7). + Do() + + verifyQuotaUser(ctx, t, client, second.Org.ID.String(), user.ID.String(), 0, 35) + verifyQuotaUser(ctx, t, client, second.Org.ID.String(), noise.ID.String(), 0, 65) + + // Workspaces owned by the user + consumed := 0 + for i := 0; i < 2; i++ { + const cost = 5 + dbfake.WorkspaceBuild(t, db, + database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: second.Org.ID, + }). + Seed(database.WorkspaceBuild{ + DailyCost: cost, + }).Do() + consumed += cost + } + + // Add some noise + // Workspace by the user in the third org + dbfake.WorkspaceBuild(t, db, + database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: third.Org.ID, + }). + Seed(database.WorkspaceBuild{ + DailyCost: 10, + }).Do() + + // Workspace by another user in third org + dbfake.WorkspaceBuild(t, db, + database.WorkspaceTable{ + OwnerID: noise.ID, + OrganizationID: third.Org.ID, + }). + Seed(database.WorkspaceBuild{ + DailyCost: 10, + }).Do() + + // Workspace by another user in second org + dbfake.WorkspaceBuild(t, db, + database.WorkspaceTable{ + OwnerID: noise.ID, + OrganizationID: second.Org.ID, + }). + Seed(database.WorkspaceBuild{ + DailyCost: 10, + }).Do() + + verifyQuotaUser(ctx, t, client, second.Org.ID.String(), user.ID.String(), consumed, 35) + }) } // nolint:paralleltest,tparallel // Tests must run serially From 5dcc74f27ebc348e89b71e02efc5d78163b98619 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Fri, 1 Nov 2024 11:43:37 -0400 Subject: [PATCH 2/5] chore: remove excess join in GetQuotaConsumedForUser query Filter is applied in original workspace query. --- coderd/database/queries.sql.go | 13 +++---------- coderd/database/queries/quotas.sql | 13 +++---------- 2 files changed, 6 insertions(+), 20 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b983d0e1bcd9d..2d3053ab60640 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6746,6 +6746,8 @@ FROM INNER JOIN workspaces on wb.workspace_id = workspaces.id WHERE + -- Only return workspaces that match the user + organization. + -- Quotas are calculated per user per organization. workspaces.owner_id = $1 AND workspaces.organization_id = $2 ORDER BY @@ -6755,16 +6757,7 @@ ORDER BY SELECT coalesce(SUM(daily_cost), 0)::BIGINT FROM - workspaces -INNER JOIN latest_builds ON - latest_builds.workspace_id = workspaces.id -WHERE - NOT deleted AND - -- We can likely remove these conditions since we check above. - -- But it does not hurt to be defensive and make sure future query changes - -- do not break anything. - workspaces.owner_id = $1 AND - workspaces.organization_id = $2 + latest_builds ` type GetQuotaConsumedForUserParams struct { diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index 7ab6189dfe8a1..709821fadbda2 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -28,6 +28,8 @@ FROM INNER JOIN workspaces on wb.workspace_id = workspaces.id WHERE + -- Only return workspaces that match the user + organization. + -- Quotas are calculated per user per organization. workspaces.owner_id = @owner_id AND workspaces.organization_id = @organization_id ORDER BY @@ -37,14 +39,5 @@ ORDER BY SELECT coalesce(SUM(daily_cost), 0)::BIGINT FROM - workspaces -INNER JOIN latest_builds ON - latest_builds.workspace_id = workspaces.id -WHERE - NOT deleted AND - -- We can likely remove these conditions since we check above. - -- But it does not hurt to be defensive and make sure future query changes - -- do not break anything. - workspaces.owner_id = @owner_id AND - workspaces.organization_id = @organization_id + latest_builds ; From 3caa9d52018c04ae27e2db684d79d3dad0ef4d7f Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Fri, 1 Nov 2024 11:46:02 -0400 Subject: [PATCH 3/5] exclude deleted workspaces' --- coderd/database/queries.sql.go | 1 + coderd/database/queries/quotas.sql | 1 + 2 files changed, 2 insertions(+) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2d3053ab60640..70157e658c8bd 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6748,6 +6748,7 @@ INNER JOIN WHERE -- Only return workspaces that match the user + organization. -- Quotas are calculated per user per organization. + NOT workspaces.deleted AND workspaces.owner_id = $1 AND workspaces.organization_id = $2 ORDER BY diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index 709821fadbda2..81da55bb24182 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -30,6 +30,7 @@ INNER JOIN WHERE -- Only return workspaces that match the user + organization. -- Quotas are calculated per user per organization. + NOT workspaces.deleted AND workspaces.owner_id = @owner_id AND workspaces.organization_id = @organization_id ORDER BY From c8eed822311aa2a6d91cece968a4fbb5bdba8059 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Fri, 1 Nov 2024 11:58:58 -0400 Subject: [PATCH 4/5] linting --- enterprise/coderd/workspacequota_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 93dd7d0d84aa7..5ec308eb6de62 100644 --- a/enterprise/coderd/workspacequota_test.go +++ b/enterprise/coderd/workspacequota_test.go @@ -323,7 +323,6 @@ func TestWorkspaceQuota(t *testing.T) { // Prepopulate database. Use dbfake as it is quicker and // easier than the api. ctx := testutil.Context(t, testutil.WaitLong) - ctx = dbauthz.AsSystemRestricted(ctx) user := dbgen.User(t, db, database.User{}) noise := dbgen.User(t, db, database.User{}) From 6420170f2de0078d22531e7837c974d402824c34 Mon Sep 17 00:00:00 2001 From: Steven Masley Date: Fri, 1 Nov 2024 12:13:12 -0400 Subject: [PATCH 5/5] build number over date --- coderd/database/queries.sql.go | 2 +- coderd/database/queries/quotas.sql | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 70157e658c8bd..1fcd1d294b0bd 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6753,7 +6753,7 @@ WHERE workspaces.organization_id = $2 ORDER BY wb.workspace_id, - wb.created_at DESC + wb.build_number DESC ) SELECT coalesce(SUM(daily_cost), 0)::BIGINT diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql index 81da55bb24182..5190057fe68bc 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -35,7 +35,7 @@ WHERE workspaces.organization_id = @organization_id ORDER BY wb.workspace_id, - wb.created_at DESC + wb.build_number DESC ) SELECT coalesce(SUM(daily_cost), 0)::BIGINT