chore: remove excess join in GetQuotaConsumedForUser query (#15338)

Emyrk · web-flow · commit 2d00b50eb6d4 · 2024-11-05T08:12:56.000-06:00
Filter is applied in original workspace query. We do not need to join
`workspaces` twice. Use build_number instead of `created_at`
for determining the last build.
diff --git 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
diff --git 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
@@ -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
