chore: scope workspace quotas to organizations (#14352)

* 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
+6 -6
View File
@@ -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) {
+8 -2
View File
@@ -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{})
+11 -5
View File
@@ -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,7 +3340,7 @@ 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
}
@@ -3348,19 +3348,25 @@ func (q *FakeQuerier) GetQuotaAllowanceForUser(_ context.Context, userID uuid.UU
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 {
+2 -2
View File
@@ -963,14 +963,14 @@ func (m metricsStore) GetProvisionerLogsAfterID(ctx context.Context, arg databas
return logs, err
}
func (m metricsStore) GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error) {
func (m metricsStore) GetQuotaAllowanceForUser(ctx context.Context, userID database.GetQuotaAllowanceForUserParams) (int64, error) {
start := time.Now()
allowance, err := m.s.GetQuotaAllowanceForUser(ctx, userID)
m.queryLatencies.WithLabelValues("GetQuotaAllowanceForUser").Observe(time.Since(start).Seconds())
return allowance, err
}
func (m metricsStore) GetQuotaConsumedForUser(ctx context.Context, ownerID uuid.UUID) (int64, error) {
func (m metricsStore) GetQuotaConsumedForUser(ctx context.Context, ownerID database.GetQuotaConsumedForUserParams) (int64, error) {
start := time.Now()
consumed, err := m.s.GetQuotaConsumedForUser(ctx, ownerID)
m.queryLatencies.WithLabelValues("GetQuotaConsumedForUser").Observe(time.Since(start).Seconds())
+2 -2
View File
@@ -1960,7 +1960,7 @@ func (mr *MockStoreMockRecorder) GetProvisionerLogsAfterID(arg0, arg1 any) *gomo
}
// GetQuotaAllowanceForUser mocks base method.
func (m *MockStore) GetQuotaAllowanceForUser(arg0 context.Context, arg1 uuid.UUID) (int64, error) {
func (m *MockStore) GetQuotaAllowanceForUser(arg0 context.Context, arg1 database.GetQuotaAllowanceForUserParams) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetQuotaAllowanceForUser", arg0, arg1)
ret0, _ := ret[0].(int64)
@@ -1975,7 +1975,7 @@ func (mr *MockStoreMockRecorder) GetQuotaAllowanceForUser(arg0, arg1 any) *gomoc
}
// GetQuotaConsumedForUser mocks base method.
func (m *MockStore) GetQuotaConsumedForUser(arg0 context.Context, arg1 uuid.UUID) (int64, error) {
func (m *MockStore) GetQuotaConsumedForUser(arg0 context.Context, arg1 database.GetQuotaConsumedForUserParams) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetQuotaConsumedForUser", arg0, arg1)
ret0, _ := ret[0].(int64)
+2 -2
View File
@@ -194,8 +194,8 @@ type sqlcQuerier interface {
GetProvisionerKeyByID(ctx context.Context, id uuid.UUID) (ProvisionerKey, error)
GetProvisionerKeyByName(ctx context.Context, arg GetProvisionerKeyByNameParams) (ProvisionerKey, error)
GetProvisionerLogsAfterID(ctx context.Context, arg GetProvisionerLogsAfterIDParams) ([]ProvisionerJobLog, error)
GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error)
GetQuotaConsumedForUser(ctx context.Context, ownerID uuid.UUID) (int64, error)
GetQuotaAllowanceForUser(ctx context.Context, arg GetQuotaAllowanceForUserParams) (int64, error)
GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error)
GetReplicaByID(ctx context.Context, id uuid.UUID) (Replica, error)
GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error)
GetTailnetAgents(ctx context.Context, id uuid.UUID) ([]TailnetAgent, error)
+8 -2
View File
@@ -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)
})
+22 -6
View File
@@ -6231,14 +6231,22 @@ 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 user_id, user_email, user_username, user_hashed_password, user_created_at, user_updated_at, user_status, user_rbac_roles, user_login_type, user_avatar_url, user_deleted, user_last_seen_at, user_quiet_hours_schedule, user_theme_preference, user_name, user_github_com_user_id, organization_id, group_name, group_id FROM group_members_expanded WHERE $1 = user_id
SELECT user_id, user_email, user_username, user_hashed_password, user_created_at, user_updated_at, user_status, user_rbac_roles, user_login_type, user_avatar_url, user_deleted, user_last_seen_at, user_quiet_hours_schedule, user_theme_preference, user_name, user_github_com_user_id, organization_id, group_name, group_id FROM group_members_expanded
WHERE
$1 = user_id AND
$2 = group_members_expanded.organization_id
) AS members
INNER JOIN groups ON
members.group_id = groups.id
`
func (q *sqlQuerier) GetQuotaAllowanceForUser(ctx context.Context, userID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, getQuotaAllowanceForUser, userID)
type GetQuotaAllowanceForUserParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
}
func (q *sqlQuerier) GetQuotaAllowanceForUser(ctx context.Context, arg GetQuotaAllowanceForUserParams) (int64, error) {
row := q.db.QueryRowContext(ctx, getQuotaAllowanceForUser, arg.UserID, arg.OrganizationID)
var column_1 int64
err := row.Scan(&column_1)
return column_1, err
@@ -6263,11 +6271,19 @@ 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 = $1 AND
workspaces.organization_id = $2
`
func (q *sqlQuerier) GetQuotaConsumedForUser(ctx context.Context, ownerID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, getQuotaConsumedForUser, ownerID)
type GetQuotaConsumedForUserParams struct {
OwnerID uuid.UUID `db:"owner_id" json:"owner_id"`
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
}
func (q *sqlQuerier) GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error) {
row := q.db.QueryRowContext(ctx, getQuotaConsumedForUser, arg.OwnerID, arg.OrganizationID)
var column_1 int64
err := row.Scan(&column_1)
return column_1, err
+9 -2
View File
@@ -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
;