Move usageAttributionId

components/dashboard/src/settings/Billing.tsx

@@ -42,9 +42,8 @@ export default function Billing() {
         if (!user) {
             return;
         }
-        const additionalData = user.additionalData || {};
-        additionalData.usageAttributionId = team ? `team:${team.id}` : `user:${user.id}`;
-        await getGitpodService().server.updateLoggedInUser({ additionalData });
+        const usageAttributionId = team ? `team:${team.id}` : `user:${user.id}`;
+        await getGitpodService().server.setUsageAttribution(usageAttributionId);
     };
 
     return (
@@ -73,9 +72,7 @@ export default function Billing() {
                         <span>Bill all my usage to:</span>
                         <DropDown
                             activeEntry={
-                                teamsWithBillingEnabled.find(
-                                    (t) => `team:${t.id}` === user?.additionalData?.usageAttributionId,
-                                )?.name
+                                teamsWithBillingEnabled.find((t) => `team:${t.id}` === user?.usageAttributionId)?.name
                             }
                             customClasses="w-32"
                             renderAsLink={true}

components/gitpod-db/src/typeorm/migration/1658394096656-UsageAttributionId.ts

@@ -0,0 +1,23 @@
+/**
+ * Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+ * Licensed under the GNU Affero General Public License (AGPL).
+ * See License-AGPL.txt in the project root for license information.
+ */
+
+import { MigrationInterface, QueryRunner } from "typeorm";
+import { columnExists } from "./helper/helper";
+
+const TABLE_NAME = "d_b_user";
+const COLUMN_NAME = "usageAttributionId";
+
+export class UsageAttributionId1658394096656 implements MigrationInterface {
+    public async up(queryRunner: QueryRunner): Promise<void> {
+        if (!(await columnExists(queryRunner, TABLE_NAME, COLUMN_NAME))) {
+            await queryRunner.query(
+                `ALTER TABLE ${TABLE_NAME} ADD COLUMN ${COLUMN_NAME} varchar(60) NOT NULL DEFAULT '', ALGORITHM=INPLACE, LOCK=NONE `,
+            );
+        }
+    }
+
+    public async down(queryRunner: QueryRunner): Promise<void> {}
+}

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

@@ -292,6 +292,7 @@ export interface GitpodServer extends JsonRpcServer<GitpodClient>, AdminServer,
     getStripePortalUrlForTeam(teamId: string): Promise<string>;
 
     listBilledUsage(attributionId: string): Promise<BillableSession[]>;
+    setUsageAttribution(usageAttribution: string): Promise<void>;
 
     /**
      * Analytics

components/gitpod-protocol/src/protocol.ts

@@ -45,6 +45,9 @@ export interface User {
     markedDeleted?: boolean;
 
     additionalData?: AdditionalUserData;
+
+    // Identifies an explicit team or user ID to which all the user's workspace usage should be attributed to (e.g. for billing purposes)
+    usageAttributionId?: string;
 }
 
 export namespace User {
@@ -199,8 +202,6 @@ export interface AdditionalUserData {
     knownGitHubOrgs?: string[];
     // Git clone URL pointing to the user's dotfile repo
     dotfileRepo?: string;
-    // Identifies an explicit team or user ID to which all the user's workspace usage should be attributed to (e.g. for billing purposes)
-    usageAttributionId?: string;
     // preferred workspace classes
     workspaceClasses?: WorkspaceClasses;
     // additional user profile data

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

@@ -1497,16 +1497,14 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
             this.stripeService.findCustomerByTeamId(teamId),
             this.userDB.findUserById(userId),
         ]);
-        if (teamCustomer && user && !user.additionalData?.usageAttributionId) {
+        if (teamCustomer && user && !user.usageAttributionId) {
             // If the user didn't explicitly choose yet where their usage should be attributed to, and
             // they join a team which accepts usage attribution (i.e. with usage-based billing enabled),
             // then we simplify the UX by automatically attributing the user's usage to that team.
             // Note: This default choice can be changed at any time by the user in their billing settings.
             const subscription = await this.stripeService.findUncancelledSubscriptionByCustomer(teamCustomer.id);
             if (subscription) {
-                user.additionalData = user.additionalData || {};
-                user.additionalData.usageAttributionId = `team:${teamId}`;
-                await this.userDB.updateUserPartial(user);
+                await this.userService.setUsageAttribution(user, AttributionId.render({ kind: "team", teamId }));
             }
         }
     }
@@ -1519,12 +1517,11 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
             await this.teamSubscription2Service.cancelTeamMemberSubscription(ts2, userId, teamMembershipId, now);
         }
         const user = await this.userDB.findUserById(userId);
-        if (user && user.additionalData?.usageAttributionId === `team:${teamId}`) {
+        if (user && user.usageAttributionId === AttributionId.render({ kind: "team", teamId })) {
             // If the user previously attributed all their usage to a given team, but they are now leaving this
             // team, then the currently selected usage attribution ID is no longer valid. In this case, we must
             // reset this ID to the default value.
-            user.additionalData.usageAttributionId = undefined;
-            await this.userDB.updateUserPartial(user);
+            await this.userService.setUsageAttribution(user, undefined);
         }
     }
 
@@ -2027,10 +2024,8 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
             await Promise.all(
                 members.map(async (m) => {
                     const u = await this.userDB.findUserById(m.userId);
-                    if (u && !u.additionalData?.usageAttributionId) {
-                        u.additionalData = u.additionalData || {};
-                        u.additionalData.usageAttributionId = `team:${teamId}`;
-                        await this.userDB.updateUserPartial(u);
+                    if (u && !u.usageAttributionId) {
+                        await this.userService.setUsageAttribution(u, `team:${teamId}`);
                     }
                 }),
             );

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

@@ -216,6 +216,7 @@ function getConfig(config: RateLimiterConfig): RateLimiterConfig {
         identifyUser: { group: "default", points: 1 },
         getIDEOptions: { group: "default", points: 1 },
         getPrebuildEvents: { group: "default", points: 1 },
+        setUsageAttribution: { group: "default", points: 1 },
     };
 
     return {

components/server/src/user/user-service.ts

@@ -16,7 +16,7 @@ import {
     WORKSPACE_TIMEOUT_EXTENDED,
     WORKSPACE_TIMEOUT_EXTENDED_ALT,
 } from "@gitpod/gitpod-protocol";
-import { CostCenterDB, ProjectDB, TermsAcceptanceDB, UserDB } from "@gitpod/gitpod-db/lib";
+import { CostCenterDB, ProjectDB, TeamDB, TermsAcceptanceDB, UserDB } from "@gitpod/gitpod-db/lib";
 import { HostContextProvider } from "../auth/host-context-provider";
 import { log } from "@gitpod/gitpod-protocol/lib/util/logging";
 import { Config } from "../config";
@@ -28,6 +28,7 @@ import { TokenService } from "./token-service";
 import { EmailAddressAlreadyTakenException, SelectAccountException } from "../auth/errors";
 import { SelectAccountPayload } from "@gitpod/gitpod-protocol/lib/auth";
 import { AttributionId } from "@gitpod/gitpod-protocol/lib/attribution";
+import { StripeService } from "../../ee/src/user/stripe-service";
 
 export interface FindUserByIdentityStrResult {
     user: User;
@@ -66,6 +67,8 @@ export class UserService {
     @inject(TermsProvider) protected readonly termsProvider: TermsProvider;
     @inject(ProjectDB) protected readonly projectDb: ProjectDB;
     @inject(CostCenterDB) protected readonly costCenterDb: CostCenterDB;
+    @inject(TeamDB) protected readonly teamDB: TeamDB;
+    @inject(StripeService) protected readonly stripeService: StripeService;
 
     /**
      * Takes strings in the form of <authHost>/<authName> and returns the matching User
@@ -226,12 +229,12 @@ export class UserService {
     async getWorkspaceUsageAttributionId(user: User, projectId?: string): Promise<string | undefined> {
         // A. Billing-based attribution
         if (this.config.enablePayment) {
-            if (!user.additionalData?.usageAttributionId) {
+            if (!user.usageAttributionId) {
                 // No explicit user attribution ID yet -- attribute all usage to the user by default (regardless of project/team).
                 return AttributionId.render({ kind: "user", userId: user.id });
             }
             // Return the user's explicit attribution ID.
-            return user.additionalData.usageAttributionId;
+            return user.usageAttributionId;
         }
 
         // B. Project-based attribution
@@ -248,6 +251,37 @@ export class UserService {
         return AttributionId.render({ kind: "team", teamId: project.teamId });
     }
 
+    async setUsageAttribution(user: User, usageAttributionId: string | undefined): Promise<void> {
+        if (typeof usageAttributionId !== "string") {
+            user.usageAttributionId = undefined;
+            await this.userDb.storeUser(user);
+            return;
+        }
+        const attributionId = AttributionId.parse(usageAttributionId);
+        if (attributionId?.kind === "user") {
+            if (user.id === attributionId.userId) {
+                user.usageAttributionId = usageAttributionId;
+                await this.userDb.storeUser(user);
+                return;
+            }
+            throw new Error("Permission denied: cannot attribute another user.");
+        }
+        if (attributionId?.kind === "team") {
+            const membership = await this.teamDB.findTeamMembership(user.id, attributionId.teamId);
+            if (!membership) {
+                throw new Error("Cannot attribute to an unrelated team.");
+            }
+            const teamCustomer = await this.stripeService.findCustomerByTeamId(attributionId.teamId);
+            if (!teamCustomer) {
+                throw new Error("Cannot attribute to team without Stripe customer.");
+            }
+            user.usageAttributionId = usageAttributionId;
+            await this.userDb.storeUser(user);
+            return;
+        }
+        throw new Error("Unexpected call arguments for `setUsageAttribution`");
+    }
+
     /**
      * This might throw `AuthException`s.
      *

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

@@ -3195,6 +3195,16 @@ export class GitpodServerImpl implements GitpodServerWithTracing, Disposable {
         throw new ResponseError(ErrorCodes.SAAS_FEATURE, `Not implemented in this version`);
     }
 
+    async setUsageAttribution(ctx: TraceContext, usageAttributionId: string): Promise<void> {
+        const user = this.checkAndBlockUser("setUsageAttribution");
+        try {
+            await this.userService.setUsageAttribution(user, usageAttributionId);
+        } catch (error) {
+            log.error("cannot set usage attribution", error, { userId: user.id, usageAttributionId });
+            throw new ResponseError(ErrorCodes.PERMISSION_DENIED, `cannot set usage attribution`);
+        }
+    }
+
     //
     //#endregion