Add getUsageLimitFor and setUsageLimitFor

Add two new methods to the server API for getting and setting usage
limits. Both new functions take an attributionId and work for both users
and teams.

For backwards compatibility, leave the `getUsageLimitForTeam` and
`setUsageLimitForTeam` methods as they are still used by the dashboard,
but change them to be implemented in terms of the more general `get/set`
methods.

Author: Andrew Farries

components/gitpod-protocol/src/gitpod-service.ts

@@ -295,6 +295,8 @@ export interface GitpodServer extends JsonRpcServer<GitpodClient>, AdminServer,
     createOrUpdateStripeCustomerForUser(currency: string): Promise<void>;
     subscribeTeamToStripe(teamId: string, setupIntentId: string): Promise<void>;
     getStripePortalUrlForTeam(teamId: string): Promise<string>;
+    getUsageLimit(attributionId: string): Promise<number | undefined>;
+    setUsageLimit(attributionId: string, usageLimit: number): Promise<void>;
     getUsageLimitForTeam(teamId: string): Promise<number | undefined>;
     setUsageLimitForTeam(teamId: string, spendingLimit: number): Promise<void>;
 

components/server/ee/src/workspace/gitpod-server-impl.ts

@@ -2146,39 +2146,75 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
         }
     }
 
-    async getUsageLimitForTeam(ctx: TraceContext, teamId: string): Promise<number | undefined> {
-        const user = this.checkAndBlockUser("getUsageLimitForTeam");
-        const team = await this.guardTeamOperation(teamId, "get");
-        await this.ensureStripeApiIsAllowed({ team });
+    async getUsageLimit(ctx: TraceContext, attributionId: string): Promise<number | undefined> {
+        const attrId = AttributionId.parse(attributionId);
+        if (attrId === undefined) {
+            log.error(`Invalid attribution id: ${attributionId}`);
+            throw new ResponseError(ErrorCodes.BAD_REQUEST, `Invalid attibution id: ${attributionId}`);
+        }
 
-        const attributionId: AttributionId = { kind: "team", teamId };
-        await this.guardCostCenterAccess(ctx, user.id, attributionId, "get");
+        const user = this.checkAndBlockUser("getUsageLimit");
+        if (attrId.kind === "team") {
+            const team = await this.guardTeamOperation(attrId.teamId, "get");
+            await this.ensureStripeApiIsAllowed({ team });
+        }
+        await this.guardCostCenterAccess(ctx, user.id, attrId, "get");
 
-        const costCenter = await this.usageServiceClientProvider.getDefault().getCostCenter(attributionId);
+        const costCenter = await this.usageServiceClientProvider.getDefault().getCostCenter(attrId);
         if (costCenter) {
             return costCenter.spendingLimit;
         }
         return undefined;
     }
 
-    async setUsageLimitForTeam(ctx: TraceContext, teamId: string, usageLimit: number): Promise<void> {
-        const user = this.checkAndBlockUser("setUsageLimitForTeam");
-        const team = await this.guardTeamOperation(teamId, "update");
-        await this.ensureStripeApiIsAllowed({ team });
+    async setUsageLimit(ctx: TraceContext, attributionId: string, usageLimit: number): Promise<void> {
+        const attrId = AttributionId.parse(attributionId);
+        if (attrId === undefined) {
+            log.error(`Invalid attribution id: ${attributionId}`);
+            throw new ResponseError(ErrorCodes.BAD_REQUEST, `Invalid attibution id: ${attributionId}`);
+        }
         if (typeof usageLimit !== "number" || usageLimit < 0) {
-            throw new ResponseError(ErrorCodes.BAD_REQUEST, "Unexpected `usageLimit` value.");
+            throw new ResponseError(ErrorCodes.BAD_REQUEST, `Unexpected usageLimit value: ${usageLimit}`);
+        }
+
+        const user = this.checkAndBlockUser("setUsageLimit");
+        switch (attrId.kind) {
+            case "team":
+                const team = await this.guardTeamOperation(attrId.teamId, "update");
+                await this.ensureStripeApiIsAllowed({ team });
+                break;
+            case "user":
+                await this.ensureStripeApiIsAllowed({ user });
+                break;
+        }
+        await this.guardCostCenterAccess(ctx, user.id, attrId, "update");
+
+        const costCenter = await this.usageServiceClientProvider.getDefault().getCostCenter(attrId);
+        if (costCenter?.billingStrategy !== "stripe") {
+            throw new ResponseError(
+                ErrorCodes.BAD_REQUEST,
+                `Setting a usage limit is not valid for non-Stripe billing strategies`,
+            );
         }
-        const attributionId: AttributionId = { kind: "team", teamId };
-        await this.guardCostCenterAccess(ctx, user.id, attributionId, "update");
 
-        const costCenter = await this.usageServiceClientProvider.getDefault().getCostCenter(attributionId);
         await this.usageServiceClientProvider.getDefault().setCostCenter({
-            id: attributionId,
+            id: attrId,
             spendingLimit: usageLimit,
             billingStrategy: costCenter?.billingStrategy || "other",
         });
     }
 
+    async getUsageLimitForTeam(ctx: TraceContext, teamId: string): Promise<number | undefined> {
+        const attributionId: AttributionId = { kind: "team", teamId: teamId };
+
+        return this.getUsageLimit(ctx, AttributionId.render(attributionId));
+    }
+
+    async setUsageLimitForTeam(ctx: TraceContext, teamId: string, usageLimit: number): Promise<void> {
+        const attributionId: AttributionId = { kind: "team", teamId: teamId };
+        return this.setUsageLimit(ctx, AttributionId.render(attributionId), usageLimit);
+    }
+
     async getNotifications(ctx: TraceContext): Promise<string[]> {
         const result = await super.getNotifications(ctx);
         const user = this.checkAndBlockUser("getNotifications");

components/server/src/auth/rate-limiter.ts

@@ -221,6 +221,8 @@ const defaultFunctions: FunctionsConfig = {
     getIDEOptions: { group: "default", points: 1 },
     getPrebuildEvents: { group: "default", points: 1 },
     setUsageAttribution: { group: "default", points: 1 },
+    getUsageLimit: { group: "default", points: 1 },
+    setUsageLimit: { group: "default", points: 1 },
     getUsageLimitForTeam: { group: "default", points: 1 },
     setUsageLimitForTeam: { group: "default", points: 1 },
     getNotifications: { group: "default", points: 1 },

components/server/src/workspace/gitpod-server-impl.ts

@@ -3238,6 +3238,14 @@ export class GitpodServerImpl implements GitpodServerWithTracing, Disposable {
         throw new ResponseError(ErrorCodes.SAAS_FEATURE, `Not implemented in this version`);
     }
 
+    async getUsageLimit(ctx: TraceContext, attributionId: string): Promise<number | undefined> {
+        throw new ResponseError(ErrorCodes.SAAS_FEATURE, `Not implemented in this version`);
+    }
+
+    async setUsageLimit(ctx: TraceContext, attributionId: string, usageLimit: number): Promise<void> {
+        throw new ResponseError(ErrorCodes.SAAS_FEATURE, `Not implemented in this version`);
+    }
+
     async getUsageLimitForTeam(ctx: TraceContext, teamId: string): Promise<number | undefined> {
         throw new ResponseError(ErrorCodes.SAAS_FEATURE, `Not implemented in this version`);
     }