[db] long running migration for filling the organizationid in workspaces

components/gitpod-db/src/container-module.ts

@@ -70,6 +70,8 @@ import { WebhookEventDBImpl } from "./typeorm/webhook-event-db-impl";
 import { PersonalAccessTokenDBImpl } from "./typeorm/personal-access-token-db-impl";
 import { UserToTeamMigrationService } from "./user-to-team-migration-service";
 import { Synchronizer } from "./typeorm/synchronizer";
+import { WorkspaceOrganizationIdMigration } from "./long-running-migration/workspace-organizationid-migration";
+import { LongRunningMigration, LongRunningMigrationService } from "./long-running-migration/long-running-migration";
 
 // THE DB container module that contains all DB implementations
 export const dbContainerModule = new ContainerModule((bind, unbind, isBound, rebind) => {
@@ -164,5 +166,9 @@ export const dbContainerModule = new ContainerModule((bind, unbind, isBound, reb
     bind(LicenseDB).to(LicenseDBImpl).inSingletonScope();
     bind(OssAllowListDB).to(OssAllowListDBImpl).inSingletonScope();
     bind(UserToTeamMigrationService).toSelf().inSingletonScope();
+    bind(WorkspaceOrganizationIdMigration).toSelf().inSingletonScope();
     bind(Synchronizer).toSelf().inSingletonScope();
+
+    bind(LongRunningMigrationService).toSelf().inSingletonScope();
+    bind(LongRunningMigration).to(WorkspaceOrganizationIdMigration).inSingletonScope();
 });

components/gitpod-db/src/long-running-migration/long-running-migration.spec.db.ts

@@ -0,0 +1,65 @@
+/**
+ * Copyright (c) 2020 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 * as chai from "chai";
+import { testContainer } from "../test-container";
+import { Synchronizer } from "../typeorm/synchronizer";
+import { TypeORM } from "../typeorm/typeorm";
+import { LongRunningMigration, LongRunningMigrationService } from "./long-running-migration";
+const expect = chai.expect;
+
+class MultiBatchesMigration implements LongRunningMigration {
+    constructor(public readonly name: string, public batches: number) {}
+
+    public getName(): string {
+        return this.name;
+    }
+
+    public async runMigrationBatch(): Promise<boolean> {
+        return --this.batches <= 0;
+    }
+}
+
+describe("long running migration service", () => {
+    const typeORM = testContainer.get<TypeORM>(TypeORM);
+
+    const wipeRepo = async () => {
+        const conn = await typeORM.getConnection();
+        await conn.query("DELETE FROM d_b_long_running_migration");
+    };
+
+    it("should migrate until completed", async () => {
+        await wipeRepo();
+        const threeBatches = new MultiBatchesMigration("threeBatches", 3);
+        const migrationService = new LongRunningMigrationService(
+            testContainer.get<TypeORM>(TypeORM),
+            testContainer.get<Synchronizer>(Synchronizer),
+            [threeBatches],
+        );
+
+        expect(await migrationService.runMigrationBatch(), "run #1").to.be.false;
+        expect(await migrationService.runMigrationBatch(), "run #2").to.be.false;
+        expect(await migrationService.runMigrationBatch(), "run #3").to.be.true;
+        expect(threeBatches.batches).to.equal(0);
+    });
+
+    it("should migrate until all are completed", async () => {
+        await wipeRepo();
+        const threeBatches = new MultiBatchesMigration("threeBatches", 3);
+        const runsOnce = new MultiBatchesMigration("runsOnce", 1);
+        const migrationService = new LongRunningMigrationService(
+            testContainer.get<TypeORM>(TypeORM),
+            testContainer.get<Synchronizer>(Synchronizer),
+            [threeBatches, runsOnce],
+        );
+
+        expect(await migrationService.runMigrationBatch(), "run #1").to.be.false;
+        expect(await migrationService.runMigrationBatch(), "run #2").to.be.false;
+        expect(await migrationService.runMigrationBatch(), "run #3").to.be.true;
+        expect(runsOnce.batches).to.equal(0);
+        expect(threeBatches.batches).to.equal(0);
+    });
+});

components/gitpod-db/src/long-running-migration/long-running-migration.ts

@@ -0,0 +1,80 @@
+/**
+ * Copyright (c) 2023 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 { log } from "@gitpod/gitpod-protocol/lib/util/logging";
+import { inject, injectable, multiInject } from "inversify";
+import { DBLongRunningMigration } from "../typeorm/entity/db-long-running-migration";
+import { Synchronizer } from "../typeorm/synchronizer";
+import { TypeORM } from "../typeorm/typeorm";
+
+export const LongRunningMigration = Symbol("LongRunningMigration");
+export interface LongRunningMigration {
+    /**
+     * The name of the migration. This is used to store the completion state in the database.
+     * Changing the name will cause the migration to be run again.
+     */
+    getName(): string;
+
+    /**
+     * Returns true if the migration has completed.
+     */
+    runMigrationBatch(): Promise<boolean>;
+}
+
+/**
+ * This service runs all registered long running migrations in batches until they are completed.
+ * It stores the completion state in the database.
+ */
+@injectable()
+export class LongRunningMigrationService {
+    public constructor(
+        @inject(TypeORM) private typeorm: TypeORM,
+        @inject(Synchronizer) private distributedLock: Synchronizer,
+        @multiInject(LongRunningMigration) private migrations: LongRunningMigration[],
+    ) {}
+
+    /**
+     * Runs a batch for all registered long running migration.
+     *
+     * @returns true if all migrations are completed.
+     */
+    async runMigrationBatch(): Promise<boolean> {
+        return this.distributedLock.synchronized("long-running-migration", "LongRunningMigrationService", async () => {
+            const repo = (await this.typeorm.getConnection()).getRepository(DBLongRunningMigration);
+            let allCompleted = true;
+            for (const migration of this.migrations) {
+                let migrationMetaData = await repo.findOne({ name: migration.getName() });
+                if (!migrationMetaData) {
+                    migrationMetaData = await repo.save({
+                        name: migration.getName(),
+                        firstRun: new Date(),
+                        lastRun: new Date(),
+                        completed: false,
+                    });
+                }
+                if (migrationMetaData.completed) {
+                    continue;
+                }
+                log.info(`Running long running migration '${migration.getName()}' ...`);
+                const now = new Date();
+                try {
+                    const completed = await migration.runMigrationBatch();
+                    log.info(
+                        `Long running migration ${migration.getName()} took ${new Date().getTime() - now.getTime()}ms`,
+                        { completed },
+                    );
+                    migrationMetaData.completed = completed;
+                } catch (e) {
+                    log.error(`Long running migration ${migration.getName()} failed`, e);
+                }
+                allCompleted = allCompleted && migrationMetaData.completed;
+                migrationMetaData.lastRun = new Date();
+                await repo.save(migrationMetaData);
+            }
+            return allCompleted;
+        });
+    }
+}

components/gitpod-db/src/long-running-migration/workspace-organizationid-migration.spec.db.ts

@@ -0,0 +1,98 @@
+/**
+ * Copyright (c) 2020 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 * as chai from "chai";
+import { testContainer } from "../test-container";
+import { TypeORM } from "../typeorm/typeorm";
+import { v4 as uuidv4 } from "uuid";
+import { WorkspaceDB } from "../workspace-db";
+import { WorkspaceOrganizationIdMigration } from "./workspace-organizationid-migration";
+const expect = chai.expect;
+
+describe("Workspace organizationid migration", () => {
+    const typeORM = testContainer.get<TypeORM>(TypeORM);
+    const migrationService = testContainer.get<WorkspaceOrganizationIdMigration>(WorkspaceOrganizationIdMigration);
+    const workspaceDB = testContainer.get<WorkspaceDB>(WorkspaceDB);
+
+    const wipeRepo = async () => {
+        const conn = await typeORM.getConnection();
+        await conn.query("DELETE FROM d_b_workspace");
+        await conn.query("DELETE FROM d_b_workspace_instance");
+    };
+
+    it("should migrate", async () => {
+        await wipeRepo();
+        const orgId = uuidv4();
+
+        const now = new Date();
+        const ws = await workspaceDB.store({
+            id: "workspace-id",
+            description: "workspace-description",
+            creationTime: now.toISOString(),
+            contextURL: "workspace-contextURL",
+            context: { title: "workspace-context-title" },
+            ownerId: uuidv4(),
+            type: "regular",
+            config: { image: "workspace-config-image" },
+        });
+        await workspaceDB.storeInstance({
+            id: "workspace-instance-id",
+            workspaceId: ws.id,
+            creationTime: now.toISOString(),
+            usageAttributionId: `team:${orgId}`,
+            ideUrl: "workspace-instance-ideUrl",
+            region: "workspace-instance-region",
+            workspaceImage: "workspace-instance-workspaceImage",
+            status: {
+                phase: "stopped",
+                conditions: {},
+            },
+        });
+
+        const nowMinus30 = new Date(now.getTime() - 24 * 60 * 60 * 1000 * 30);
+        const ws1 = await workspaceDB.store({
+            id: "workspace-id1",
+            description: "workspace-description",
+            creationTime: nowMinus30.toISOString(),
+            contextURL: "workspace-contextURL",
+            context: { title: "workspace-context-title" },
+            ownerId: uuidv4(),
+            type: "regular",
+            config: { image: "workspace-config-image" },
+        });
+        await workspaceDB.storeInstance({
+            id: "workspace-instance-id1",
+            workspaceId: ws1.id,
+            creationTime: nowMinus30.toISOString(),
+            usageAttributionId: `team:${orgId}`,
+            ideUrl: "workspace-instance-ideUrl",
+            region: "workspace-instance-region",
+            workspaceImage: "workspace-instance-workspaceImage",
+            status: {
+                phase: "stopped",
+                conditions: {},
+            },
+        });
+
+        const conn = await typeORM.getConnection();
+        const getMigrationCount = async () => {
+            const result = await conn.query(
+                `SELECT count(*) as migrated FROM d_b_workspace WHERE organizationId='${orgId}'`,
+            );
+            return Number.parseInt(result[0].migrated);
+        };
+
+        expect(await getMigrationCount()).to.be.eq(0);
+
+        await migrationService.runMigrationBatch();
+
+        expect(await getMigrationCount()).to.be.eq(1);
+
+        await migrationService.runMigrationBatch();
+
+        expect(await getMigrationCount()).to.be.eq(2);
+    });
+});

components/gitpod-db/src/long-running-migration/workspace-organizationid-migration.ts

@@ -0,0 +1,68 @@
+/**
+ * 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 { log } from "@gitpod/gitpod-protocol/lib/util/logging";
+import { inject, injectable } from "inversify";
+import { TypeORM } from "../typeorm/typeorm";
+import { LongRunningMigration } from "./long-running-migration";
+
+const BATCH_OFFSET = 5 * 24 * 60 * 60 * 1000; /* 5 days */
+
+@injectable()
+export class WorkspaceOrganizationIdMigration implements LongRunningMigration {
+    @inject(TypeORM) protected readonly typeorm: TypeORM;
+
+    getName(): string {
+        return "WorkspaceOrganizationIdMigration";
+    }
+
+    public async runMigrationBatch(): Promise<boolean> {
+        const conn = await this.typeorm.getConnection();
+        const runner = conn.createQueryRunner();
+        const minCreationTimeMigrated = (
+            await runner.query(
+                "SELECT min(creationTime) as minTime FROM d_b_workspace WHERE organizationId IS NOT NULL",
+            )
+        )[0].minTime;
+        const minCreationTimeTotal = (await runner.query("SELECT min(creationTime) as minTime FROM d_b_workspace"))[0]
+            .minTime;
+
+        let endDate = minCreationTimeMigrated ? new Date(Date.parse(minCreationTimeMigrated)) : new Date();
+        let startDate = new Date(endDate.getTime() - BATCH_OFFSET);
+        log.info(`Running migration with start date: ${startDate}, end date: ${endDate}`);
+
+        let result;
+        do {
+            const query = `
+            UPDATE d_b_workspace w
+            JOIN (
+                SELECT workspaceId, MAX(creationTime) as maxCreationTime
+                FROM d_b_workspace_instance
+                WHERE usageAttributionId LIKE 'team:%'
+                GROUP BY workspaceId
+                ) wi ON w.id = wi.workspaceId
+                SET w.organizationid = (
+                    SELECT substr(usageAttributionId, 6)
+                    FROM d_b_workspace_instance
+                    WHERE workspaceId = wi.workspaceId AND creationTime = wi.maxCreationTime
+                    )
+                    WHERE
+                    w.creationTime >= '${startDate.toISOString()}' and
+                    w.creationTime < '${endDate.toISOString()}' and
+                    w.organizationId IS NULL and
+                    w.softDeleted IS NULL
+            `;
+            result = await runner.query(query);
+            log.info(`Migrated ${result.affectedRows} workspaces. Start date: ${startDate}, end date: ${endDate}`, {
+                query,
+            });
+            endDate = new Date(endDate.getTime() - BATCH_OFFSET);
+            startDate = new Date(startDate.getTime() - BATCH_OFFSET);
+        } while (result.affectedRows === 0 && endDate.toISOString() > minCreationTimeTotal);
+        const completed = result.affectedRows === 0;
+        return completed;
+    }
+}

components/gitpod-db/src/typeorm/entity/db-long-running-migration.ts

@@ -0,0 +1,22 @@
+/**
+ * Copyright (c) 2023 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 { Entity, Column, PrimaryColumn } from "typeorm";
+
+@Entity()
+export class DBLongRunningMigration {
+    @PrimaryColumn()
+    name: string;
+
+    @Column({ type: "timestamp", precision: 6 })
+    firstRun: Date;
+
+    @Column({ type: "timestamp", precision: 6 })
+    lastRun: Date;
+
+    @Column()
+    completed: boolean;
+}

components/gitpod-db/src/typeorm/migration/1675863598550-LongRunningMigration.ts

@@ -0,0 +1,29 @@
+/**
+ * Copyright (c) 2023 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 { tableExists } from "./helper/helper";
+
+export class LongRunningMigration1675863598550 implements MigrationInterface {
+    public async up(queryRunner: QueryRunner): Promise<void> {
+        if (!(await tableExists(queryRunner, "d_b_long_running_migration"))) {
+            await queryRunner.query(
+                `
+                    CREATE TABLE IF NOT EXISTS d_b_long_running_migration (
+                        name varchar(255) NOT NULL,
+                        firstRun timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
+                        lastRun timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
+                        completed tinyint(4) NOT NULL DEFAULT '0',
+                        _lastModified timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
+                        PRIMARY KEY (name)
+                    )
+                `,
+            );
+        }
+    }
+
+    public async down(queryRunner: QueryRunner): Promise<void> {}
+}