coder · Emyrk · Aug 21, 2024 · Aug 19, 2024 · Aug 19, 2024 · Aug 19, 2024
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -1823,20 +1823,20 @@ func (q *querier) GetProvisionerLogsAfterID(ctx context.Context, arg database.Ge
 	return q.db.GetProvisionerLogsAfterID(ctx, arg)
 }
 
-func (q *querier) GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error) {
-	err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUserObject(userID))
+func (q *querier) GetQuotaAllowanceForUser(ctx context.Context, params database.GetQuotaAllowanceForUserParams) (int64, error) {
+	err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUserObject(params.UserID))
 	if err != nil {
 		return -1, err
 	}
-	return q.db.GetQuotaAllowanceForUser(ctx, userID)
+	return q.db.GetQuotaAllowanceForUser(ctx, params)
 }
 
-func (q *querier) GetQuotaConsumedForUser(ctx context.Context, userID uuid.UUID) (int64, error) {
-	err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUserObject(userID))
+func (q *querier) GetQuotaConsumedForUser(ctx context.Context, params database.GetQuotaConsumedForUserParams) (int64, error) {
+	err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUserObject(params.OwnerID))
 	if err != nil {
 		return -1, err
 	}
-	return q.db.GetQuotaConsumedForUser(ctx, userID)
+	return q.db.GetQuotaConsumedForUser(ctx, params)
 }
 
 func (q *querier) GetReplicaByID(ctx context.Context, id uuid.UUID) (database.Replica, error) {

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1075,11 +1075,17 @@ func (s *MethodTestSuite) TestUser() {
 	}))
 	s.Run("GetQuotaAllowanceForUser", s.Subtest(func(db database.Store, check *expects) {
 		u := dbgen.User(s.T(), db, database.User{})
-		check.Args(u.ID).Asserts(u, policy.ActionRead).Returns(int64(0))
+		check.Args(database.GetQuotaAllowanceForUserParams{
+			UserID:         u.ID,
+			OrganizationID: uuid.New(),
+		}).Asserts(u, policy.ActionRead).Returns(int64(0))
 	}))
 	s.Run("GetQuotaConsumedForUser", s.Subtest(func(db database.Store, check *expects) {
 		u := dbgen.User(s.T(), db, database.User{})
-		check.Args(u.ID).Asserts(u, policy.ActionRead).Returns(int64(0))
+		check.Args(database.GetQuotaConsumedForUserParams{
+			OwnerID:        u.ID,
+			OrganizationID: uuid.New(),
+		}).Asserts(u, policy.ActionRead).Returns(int64(0))
 	}))
 	s.Run("GetUserByEmailOrUsername", s.Subtest(func(db database.Store, check *expects) {
 		u := dbgen.User(s.T(), db, database.User{})

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -3314,13 +3314,13 @@ func (q *FakeQuerier) GetProvisionerLogsAfterID(_ context.Context, arg database.
 	return logs, nil
 }
 
-func (q *FakeQuerier) GetQuotaAllowanceForUser(_ context.Context, userID uuid.UUID) (int64, error) {
+func (q *FakeQuerier) GetQuotaAllowanceForUser(_ context.Context, params database.GetQuotaAllowanceForUserParams) (int64, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
 	var sum int64
 	for _, member := range q.groupMembers {
-		if member.UserID != userID {
+		if member.UserID != params.UserID {
 			continue
 		}
 		if _, err := q.getOrganizationByIDNoLock(member.GroupID); err == nil {
@@ -3340,27 +3340,33 @@ func (q *FakeQuerier) GetQuotaAllowanceForUser(_ context.Context, userID uuid.UU
 	// Grab the quota for the Everyone group iff the user is a member of
 	// said organization.
 	for _, mem := range q.organizationMembers {
-		if mem.UserID != userID {
+		if mem.UserID != params.UserID {
 			continue
 		}
 
 		group, err := q.getGroupByIDNoLock(context.Background(), mem.OrganizationID)
 		if err != nil {
 			return -1, xerrors.Errorf("failed to get everyone group for org %q", mem.OrganizationID.String())
 		}
+		if group.OrganizationID != params.OrganizationID {
+			continue
+		}
 		sum += int64(group.QuotaAllowance)
 	}
 
 	return sum, nil
 }
 
-func (q *FakeQuerier) GetQuotaConsumedForUser(_ context.Context, userID uuid.UUID) (int64, error) {
+func (q *FakeQuerier) GetQuotaConsumedForUser(_ context.Context, params database.GetQuotaConsumedForUserParams) (int64, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
 	var sum int64
 	for _, workspace := range q.workspaces {
-		if workspace.OwnerID != userID {
+		if workspace.OwnerID != params.OwnerID {
+			continue
+		}
+		if workspace.OrganizationID != params.OrganizationID {
 			continue
 		}
 		if workspace.Deleted {

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
@@ -604,7 +604,10 @@ func TestWorkspaceQuotas(t *testing.T) {
 			db2sdk.List([]database.OrganizationMember{memOne, memTwo}, orgMemberIDs))
 
 		// Check the quota is correct.
-		allowance, err := db.GetQuotaAllowanceForUser(ctx, one.ID)
+		allowance, err := db.GetQuotaAllowanceForUser(ctx, database.GetQuotaAllowanceForUserParams{
+			UserID:         one.ID,
+			OrganizationID: org.ID,
+		})
 		require.NoError(t, err)
 		require.Equal(t, int64(50), allowance)
 
@@ -617,7 +620,10 @@ func TestWorkspaceQuotas(t *testing.T) {
 		require.NoError(t, err)
 
 		// Ensure allowance remains the same
-		allowance, err = db.GetQuotaAllowanceForUser(ctx, one.ID)
+		allowance, err = db.GetQuotaAllowanceForUser(ctx, database.GetQuotaAllowanceForUserParams{
+			UserID:         one.ID,
+			OrganizationID: org.ID,
+		})
 		require.NoError(t, err)
 		require.Equal(t, int64(50), allowance)
 	})

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
@@ -5,7 +5,10 @@ FROM
 	(
 		-- Select all groups this user is a member of. This will also include
 		-- the "Everyone" group for organizations the user is a member of.
-		SELECT * FROM group_members_expanded WHERE @user_id = user_id
+		SELECT * FROM group_members_expanded
+		         WHERE
+		             @user_id = user_id AND
+		             @organization_id = group_members_expanded.organization_id
 	) AS members
 INNER JOIN groups ON
 	members.group_id = groups.id
@@ -30,4 +33,8 @@ FROM
 	workspaces
 JOIN latest_builds ON
 	latest_builds.workspace_id = workspaces.id
-WHERE NOT deleted AND workspaces.owner_id = $1;
+WHERE NOT
+	deleted AND
+	workspaces.owner_id = @owner_id AND
+	workspaces.organization_id = @organization_id
+;
diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go
@@ -572,8 +572,8 @@ type WorkspaceQuota struct {
 	Budget          int `json:"budget"`
 }
 
-func (c *Client) WorkspaceQuota(ctx context.Context, userID string) (WorkspaceQuota, error) {
-	res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil)
+func (c *Client) WorkspaceQuota(ctx context.Context, organizationID string, userID string) (WorkspaceQuota, error) {
+	res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/organizations/%s/members/%s/workspace-quota", organizationID, userID), nil)
 	if err != nil {
 		return WorkspaceQuota{}, err
 	}