opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: scope workspace quotas to organizations (#14352)

Browse Source 
* chore: scope workspace quotas to organizations

Quotas are now a function of (user_id, organization_id). They are
still sourced from groups. Deprecate the old api endpoint.
This commit is contained in:
Steven Masley
2024-08-21 09:25:20 -05:00
committed by GitHub
parent fa733318e0
commit a359879af5
16 changed files with 309 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/coderd.go
+14 -1
View File
@@ -275,6 +275,19 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
r.Delete("/organizations/{organization}/members/roles/{roleName}", api.deleteOrgRole)
})
r.Group(func(r chi.Router) {
r.Use(
apiKeyMiddleware,
httpmw.ExtractOrganizationParam(api.Database),
// Intentionally using ExtractUser instead of ExtractMember.
// It is possible for a member to be removed from an org, in which
// case their orphaned workspaces still exist. We only need
// the user_id for the query.
httpmw.ExtractUserParam(api.Database),
)
r.Get("/organizations/{organization}/members/{user}/workspace-quota", api.workspaceQuota)
})
r.Route("/organizations/{organization}/groups", func(r chi.Router) {
r.Use(
apiKeyMiddleware,
@@ -365,7 +378,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
)
r.Route("/{user}", func(r chi.Router) {
r.Use(httpmw.ExtractUserParam(options.Database))
r.Get("/", api.workspaceQuota)
r.Get("/", api.workspaceQuotaByUser)
})
})
r.Route("/appearance", func(r chi.Router) {
enterprise/coderd/workspacequota.go
+45 -12
View File
@@ -6,6 +6,7 @@ import (
"errors"
"net/http"
"github.com/go-chi/chi/v5"
"github.com/google/uuid"
"cdr.dev/slog"
@@ -13,7 +14,6 @@ import (
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/coderd/rbac/policy"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/provisionerd/proto"
)
@@ -48,12 +48,18 @@ func (c *committer) CommitQuota(
)
err = c.Database.InTx(func(s database.Store) error {
var err error
consumed, err = s.GetQuotaConsumedForUser(ctx, workspace.OwnerID)
consumed, err = s.GetQuotaConsumedForUser(ctx, database.GetQuotaConsumedForUserParams{
OwnerID: workspace.OwnerID,
OrganizationID: workspace.OrganizationID,
})
if err != nil {
return err
}
budget, err = s.GetQuotaAllowanceForUser(ctx, workspace.OwnerID)
budget, err = s.GetQuotaAllowanceForUser(ctx, database.GetQuotaAllowanceForUserParams{
UserID: workspace.OwnerID,
OrganizationID: workspace.OrganizationID,
})
if err != nil {
return err
}
@@ -112,22 +118,43 @@ func (c *committer) CommitQuota(
}, nil
}
// @Summary Get workspace quota by user
// @ID get-workspace-quota-by-user
// @Summary Get workspace quota by user deprecated
// @ID get-workspace-quota-by-user-deprecated
// @Security CoderSessionToken
// @Produce json
// @Tags Enterprise
// @Param user path string true "User ID, name, or me"
// @Success 200 {object} codersdk.WorkspaceQuota
// @Router /workspace-quota/{user} [get]
func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
user := httpmw.UserParam(r)
if !api.AGPL.Authorize(r, policy.ActionRead, user) {
httpapi.ResourceNotFound(rw)
// @Deprecated this endpoint will be removed, use /organizations/{organization}/members/{user}/workspace-quota instead
func (api *API) workspaceQuotaByUser(rw http.ResponseWriter, r *http.Request) {
defaultOrg, err := api.Database.GetDefaultOrganization(r.Context())
if err != nil {
httpapi.InternalServerError(rw, err)
return
}
// defer to the new endpoint using default org as the organization
chi.RouteContext(r.Context()).URLParams.Add("organization", defaultOrg.ID.String())
mw := httpmw.ExtractOrganizationParam(api.Database)
mw(http.HandlerFunc(api.workspaceQuota)).ServeHTTP(rw, r)
}
// @Summary Get workspace quota by user
// @ID get-workspace-quota-by-user
// @Security CoderSessionToken
// @Produce json
// @Tags Enterprise
// @Param user path string true "User ID, name, or me"
// @Param organization path string true "Organization ID" format(uuid)
// @Success 200 {object} codersdk.WorkspaceQuota
// @Router /organizations/{organization}/members/{user}/workspace-quota [get]
func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
var (
organization = httpmw.OrganizationParam(r)
user = httpmw.UserParam(r)
)
api.entitlementsMu.RLock()
licensed := api.entitlements.Features[codersdk.FeatureTemplateRBAC].Enabled
api.entitlementsMu.RUnlock()
@@ -136,7 +163,10 @@ func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
var quotaAllowance int64 = -1
if licensed {
var err error
quotaAllowance, err = api.Database.GetQuotaAllowanceForUser(r.Context(), user.ID)
quotaAllowance, err = api.Database.GetQuotaAllowanceForUser(r.Context(), database.GetQuotaAllowanceForUserParams{
UserID: user.ID,
OrganizationID: organization.ID,
})
if err != nil {
httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to get allowance",
@@ -146,7 +176,10 @@ func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
}
}
quotaConsumed, err := api.Database.GetQuotaConsumedForUser(r.Context(), user.ID)
quotaConsumed, err := api.Database.GetQuotaConsumedForUser(r.Context(), database.GetQuotaConsumedForUserParams{
OwnerID: user.ID,
OrganizationID: organization.ID,
})
if err != nil {
httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to get consumed",
enterprise/coderd/workspacequota_test.go
+52 -19
View File
@@ -2,6 +2,9 @@ package coderd_test
import (
"context"
"encoding/json"
"fmt"
"net/http"
"sync"
"testing"
@@ -20,15 +23,31 @@ import (
"github.com/coder/coder/v2/testutil"
)
func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, consumed, total int) {
func verifyQuota(ctx context.Context, t *testing.T, client *codersdk.Client, organizationID string, consumed, total int) {
t.Helper()
got, err := client.WorkspaceQuota(ctx, codersdk.Me)
got, err := client.WorkspaceQuota(ctx, organizationID, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, codersdk.WorkspaceQuota{
Budget: total,
CreditsConsumed: consumed,
}, got)
// 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)
require.NoError(t, err, "deprecated endpoint")
// Only continue to check if the values differ
if deprecatedGot.Budget != got.Budget || deprecatedGot.CreditsConsumed != got.CreditsConsumed {
org, err := client.OrganizationByName(ctx, organizationID)
if err != nil {
return
}
if org.IsDefault {
require.Equal(t, got, deprecatedGot)
}
}
}
func TestWorkspaceQuota(t *testing.T) {
@@ -52,14 +71,14 @@ func TestWorkspaceQuota(t *testing.T) {
})
coderdtest.NewProvisionerDaemon(t, api.AGPL)
verifyQuota(ctx, t, client, 0, 0)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0)
// Patch the 'Everyone' group to verify its quota allowance is being accounted for.
_, err := client.PatchGroup(ctx, user.OrganizationID, codersdk.PatchGroupRequest{
QuotaAllowance: ptr.Ref(1),
})
require.NoError(t, err)
verifyQuota(ctx, t, client, 0, 1)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 1)
// Add user to two groups, granting them a total budget of 4.
group1, err := client.CreateGroup(ctx, user.OrganizationID, codersdk.CreateGroupRequest{
@@ -84,7 +103,7 @@ func TestWorkspaceQuota(t *testing.T) {
})
require.NoError(t, err)
verifyQuota(ctx, t, client, 0, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 4)
authToken := uuid.NewString()
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
@@ -123,14 +142,14 @@ func TestWorkspaceQuota(t *testing.T) {
}()
}
wg.Wait()
verifyQuota(ctx, t, client, 4, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
// Next one must fail
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
// Consumed shouldn't bump
verifyQuota(ctx, t, client, 4, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
require.Contains(t, build.Job.Error, "quota")
@@ -146,7 +165,7 @@ func TestWorkspaceQuota(t *testing.T) {
})
require.NoError(t, err)
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
verifyQuota(ctx, t, client, 3, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 3, 4)
break
}
@@ -154,7 +173,7 @@ func TestWorkspaceQuota(t *testing.T) {
workspace = coderdtest.CreateWorkspace(t, client, template.ID)
build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
verifyQuota(ctx, t, client, 4, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
})
@@ -174,14 +193,14 @@ func TestWorkspaceQuota(t *testing.T) {
})
coderdtest.NewProvisionerDaemon(t, api.AGPL)
verifyQuota(ctx, t, client, 0, 0)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 0)
// Patch the 'Everyone' group to verify its quota allowance is being accounted for.
_, err := client.PatchGroup(ctx, user.OrganizationID, codersdk.PatchGroupRequest{
QuotaAllowance: ptr.Ref(4),
})
require.NoError(t, err)
verifyQuota(ctx, t, client, 0, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 0, 4)
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
@@ -208,7 +227,7 @@ func TestWorkspaceQuota(t *testing.T) {
assert.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
}
wg.Wait()
verifyQuota(ctx, t, client, 4, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
// Next one must fail
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
@@ -216,21 +235,21 @@ func TestWorkspaceQuota(t *testing.T) {
require.Contains(t, build.Job.Error, "quota")
// Consumed shouldn't bump
verifyQuota(ctx, t, client, 4, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
require.Equal(t, codersdk.WorkspaceStatusFailed, build.Status)
build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStop)
build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
// Quota goes down one
verifyQuota(ctx, t, client, 3, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 3, 4)
require.Equal(t, codersdk.WorkspaceStatusStopped, build.Status)
build = coderdtest.CreateWorkspaceBuild(t, client, workspaces[0], database.WorkspaceTransitionStart)
build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
// Quota goes back up
verifyQuota(ctx, t, client, 4, 4)
verifyQuota(ctx, t, client, user.OrganizationID.String(), 4, 4)
require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
})
@@ -273,13 +292,27 @@ func TestWorkspaceQuota(t *testing.T) {
})
require.NoError(t, err)
verifyQuota(ctx, t, member, 0, 30)
// This currently reports the total site wide quotas. We might want to
// org scope this api call in the future.
verifyQuota(ctx, t, owner, 0, 45)
verifyQuota(ctx, t, member, first.OrganizationID.String(), 0, 30)
// Verify org scoped quota limits
verifyQuota(ctx, t, owner, first.OrganizationID.String(), 0, 30)
verifyQuota(ctx, t, owner, second.ID.String(), 0, 15)
})
}
func deprecatedQuotaEndpoint(ctx context.Context, client *codersdk.Client, userID string) (codersdk.WorkspaceQuota, error) {
res, err := client.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil)
if err != nil {
return codersdk.WorkspaceQuota{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return codersdk.WorkspaceQuota{}, codersdk.ReadBodyAsError(res)
}
var quota codersdk.WorkspaceQuota
return quota, json.NewDecoder(res.Body).Decode(&quota)
}
func planWithCost(cost int32) []*proto.Response {
return []*proto.Response{{
Type: &proto.Response_Plan{