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/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b983d0e1bcd9d..1fcd1d294b0bd 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6746,25 +6746,19 @@ 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. + NOT workspaces.deleted AND workspaces.owner_id = $1 AND workspaces.organization_id = $2 ORDER BY wb.workspace_id, - wb.created_at DESC + wb.build_number DESC ) 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..5190057fe68bc 100644 --- a/coderd/database/queries/quotas.sql +++ b/coderd/database/queries/quotas.sql @@ -28,23 +28,17 @@ 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. + NOT workspaces.deleted AND workspaces.owner_id = @owner_id AND 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 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 ; diff --git a/enterprise/coderd/workspacequota_test.go b/enterprise/coderd/workspacequota_test.go index 13142f11e5717..5ec308eb6de62 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,95 @@ 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) + + 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