[server] migrate ws without usageattribution (#17485)

Author: svenefftinge

components/gitpod-db/src/user-to-team-migration-service.spec.db.ts

@@ -158,4 +158,41 @@ describe("Migration Service", () => {
         const teams = await teamDB.findTeamsByUser(user.id);
         expect(teams[0].name).to.be.eq("X Organization");
     });
+
+    it("should update 'organizationId' for workspace without attributionId", async () => {
+        await wipeRepo();
+        const user = await userDB.newUser();
+        await userDB.storeUser(user);
+
+        const ws = await workspaceDB.store({
+            id: uuidv4(),
+            creationTime: new Date().toISOString(),
+            ownerId: user.id,
+            config: {},
+            context: {
+                title: "test",
+            },
+            contextURL: "https://gitpod.io",
+            type: "regular",
+            description: "test",
+        });
+
+        await workspaceDB.storeInstance({
+            id: uuidv4(),
+            creationTime: new Date().toISOString(),
+            region: "eu-west-1",
+            ideUrl: "https://ide.eu-west-1.aws.com",
+            workspaceImage: "test",
+            status: {
+                conditions: {},
+                phase: "stopped",
+            },
+            workspaceId: ws.id,
+        });
+
+        await migrationService.migrateUser(user);
+        const wsAndI = await workspaceDB.findWorkspaceAndInstance(ws.id);
+        const teams = await teamDB.findTeamsByUser(user.id);
+        expect(wsAndI?.organizationId).to.be.eq(teams[0].id);
+    });
 });

components/gitpod-db/src/user-to-team-migration-service.ts

@@ -4,7 +4,7 @@
  * See License.AGPL.txt in the project root for license information.
  */
 
-import { AdditionalUserData, Team, User } from "@gitpod/gitpod-protocol";
+import { AdditionalUserData, Team, User, WorkspaceInfo } from "@gitpod/gitpod-protocol";
 import { ErrorCodes } from "@gitpod/gitpod-protocol/lib/messaging/error";
 import { inject, injectable } from "inversify";
 import { ProjectDB } from "./project-db";
@@ -15,6 +15,7 @@ import { TypeORM } from "./typeorm/typeorm";
 import { log, LogContext } from "@gitpod/gitpod-protocol/lib/util/logging";
 import { UserDB } from "./user-db";
 import { Synchronizer } from "./typeorm/synchronizer";
+import { AttributionId } from "@gitpod/gitpod-protocol/lib/attribution";
 
 @injectable()
 export class UserToTeamMigrationService {
@@ -135,6 +136,13 @@ export class UserToTeamMigrationService {
         );
         log.info(ctx, "Migrated workspaces.", { teamId: team.id, result });
 
+        // Ensure there are no workspaces without an organizationId. This is necessary because very old workspace instance don't have an attributionId.
+        const workspaces = await this.workspaceDB.find({
+            userId: user.id,
+        });
+        await this.updateWorkspacesOrganizationId(workspaces, team.id);
+        log.info(ctx, "Updated workspaces.", { teamId: team.id });
+
         result = await conn.query("UPDATE d_b_usage SET attributionId = ? WHERE attributionId = ?", [
             newAttribution,
             oldAttribution,
@@ -153,4 +161,30 @@ export class UserToTeamMigrationService {
         const teams = await this.teamDB.findTeamsByUser(user.id);
         return teams.length === 0 || !user.additionalData?.isMigratedToTeamOnlyAttribution;
     }
+
+    async updateWorkspacesOrganizationId(workspaces: WorkspaceInfo[], userOrgId: string): Promise<WorkspaceInfo[]> {
+        return await Promise.all(
+            workspaces.map(async (ws) => {
+                if (!ws.workspace.organizationId) {
+                    const attrId =
+                        ws.latestInstance?.usageAttributionId &&
+                        AttributionId.parse(ws.latestInstance.usageAttributionId);
+                    if (attrId && attrId.kind === "team") {
+                        ws.workspace.organizationId = attrId.teamId;
+                    } else {
+                        ws.workspace.organizationId = userOrgId;
+                    }
+                    await this.workspaceDB.updatePartial(ws.workspace.id, {
+                        organizationId: ws.workspace.organizationId,
+                    });
+                }
+                return ws;
+            }),
+        );
+    }
+
+    async getUserOrganization(user: User): Promise<Team> {
+        const teams = await this.teamDB.findTeamsByUser(user.id);
+        return teams.find((t) => t.name === user.name || t.name === user.fullName) || teams[0];
+    }
 }

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

@@ -1206,6 +1206,17 @@ export class GitpodServerImpl implements GitpodServerWithTracing, Disposable {
                 ),
             ),
         );
+
+        // we need to update old workspaces on the fly that didn't get an orgId because we lack attribution on their instances.
+        // this can be removed eventually.
+        if (user.additionalData?.isMigratedToTeamOnlyAttribution) {
+            try {
+                const userOrg = await this.userToTeamMigrationService.getUserOrganization(user);
+                await this.userToTeamMigrationService.updateWorkspacesOrganizationId(res, userOrg.id);
+            } catch (error) {
+                log.error({ userId: user.id }, "Error updating workspaces without orgId.", error);
+            }
+        }
         return res;
     }
 
@@ -1292,11 +1303,24 @@ export class GitpodServerImpl implements GitpodServerWithTracing, Disposable {
     }
 
     protected async internalGetWorkspace(id: string, db: WorkspaceDB): Promise<Workspace> {
-        const ws = await db.findById(id);
-        if (!ws) {
+        const workspace = await db.findById(id);
+        if (!workspace) {
             throw new ResponseError(ErrorCodes.NOT_FOUND, "Workspace not found.");
         }
-        return ws;
+        if (!workspace.organizationId && this.user?.additionalData?.isMigratedToTeamOnlyAttribution) {
+            try {
+                log.info({ userId: this.user.id }, "Updating workspace without orgId.");
+                const userOrg = await this.userToTeamMigrationService.getUserOrganization(this.user);
+                const latestInstance = await this.workspaceDb.trace({}).findCurrentInstance(workspace.id);
+                this.userToTeamMigrationService.updateWorkspacesOrganizationId(
+                    [{ workspace, latestInstance }],
+                    userOrg.id,
+                );
+            } catch (error) {
+                log.error({ userId: this.user.id }, "Error updating workspaces without orgId.", error);
+            }
+        }
+        return workspace;
     }
 
     private async findRunningInstancesForContext(