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.
gitpod-io · Sep 14, 2022 · 6d9c770 · 6d9c770
1 parent 0c63b14
commit 6d9c770
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 18 deletions.
diff --git a/components/gitpod-protocol/src/gitpod-service.ts b/components/gitpod-protocol/src/gitpod-service.ts
@@ -293,6 +293,8 @@ export interface GitpodServer extends JsonRpcServer<GitpodClient>, AdminServer,
     subscribeTeamToStripe(teamId: string, setupIntentId: string): Promise<void>;
     getStripePortalUrl(attributionId: string): Promise<string>;
     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>;
 

diff --git a/components/server/ee/src/workspace/gitpod-server-impl.ts b/components/server/ee/src/workspace/gitpod-server-impl.ts
@@ -2218,44 +2218,73 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
         return this.getStripePortalUrl(ctx, AttributionId.render(attributionId));
     }
 
-    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");
+        await this.guardCostCenterAccess(ctx, user.id, attrId, "get");
 
-        const costCenter = await this.usageService.getCostCenter({
-            attributionId: AttributionId.render(attributionId),
-        });
+        const costCenter = await this.usageService.getCostCenter({ attributionId });
         if (costCenter?.costCenter) {
             return costCenter.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;
         }
-        const attrId: AttributionId = { kind: "team", teamId };
         await this.guardCostCenterAccess(ctx, user.id, attrId, "update");
-        const attributionId = AttributionId.render(attrId);
-        const costCenter = await this.usageService.getCostCenter({ attributionId });
+
+        const response = await this.usageService.getCostCenter({ attributionId });
+        if (response?.costCenter?.billingStrategy !== CostCenter_BillingStrategy.BILLING_STRATEGY_STRIPE) {
+            throw new ResponseError(
+                ErrorCodes.BAD_REQUEST,
+                `Setting a usage limit is not valid for non-Stripe billing strategies`,
+            );
+        }
         await this.usageService.setCostCenter({
             costCenter: {
                 attributionId,
                 spendingLimit: usageLimit,
                 billingStrategy:
-                    costCenter?.costCenter?.billingStrategy || CostCenter_BillingStrategy.BILLING_STRATEGY_OTHER,
+                    response?.costCenter?.billingStrategy || CostCenter_BillingStrategy.BILLING_STRATEGY_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");

diff --git a/components/server/src/auth/rate-limiter.ts b/components/server/src/auth/rate-limiter.ts
@@ -219,6 +219,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 },

diff --git a/components/server/src/workspace/gitpod-server-impl.ts b/components/server/src/workspace/gitpod-server-impl.ts
@@ -3158,6 +3158,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`);
     }