pangolin/server/setup/scriptsPg/1.10.0.ts

import { db } from "@server/db/pg/driver";
import { sql } from "drizzle-orm";
import { __DIRNAME, APP_PATH } from "@server/lib/consts";
import { readFileSync } from "fs";
import path, { join } from "path";

const version = "1.10.0";

export default async function migration() {
    console.log(`Running setup script ${version}...`);

    try {
        const resources = await db.execute(sql`
            SELECT "resourceId" FROM "resources"
        `);

        const siteResources = await db.execute(sql`
            SELECT "siteResourceId" FROM "siteResources"
        `);

        await db.execute(sql`BEGIN`);

        await db.execute(
            sql`ALTER TABLE "exitNodes" ADD COLUMN "region" text;`
        );

        await db.execute(
            sql`ALTER TABLE "idpOidcConfig" ADD COLUMN "variant" text DEFAULT 'oidc' NOT NULL;`
        );

        await db.execute(
            sql`ALTER TABLE "resources" ADD COLUMN "niceId" text DEFAULT '' NOT NULL;`
        );

        await db.execute(
            sql`ALTER TABLE "siteResources" ADD COLUMN "niceId" text DEFAULT '' NOT NULL;`
        );

        await db.execute(
            sql`ALTER TABLE "userOrgs" ADD COLUMN "autoProvisioned" boolean DEFAULT false;`
        );

        await db.execute(
            sql`ALTER TABLE "targets" ADD COLUMN "pathMatchType" text;`
        );

        await db.execute(sql`ALTER TABLE "targets" ADD COLUMN "path" text;`);

        await db.execute(
            sql`ALTER TABLE "resources" ADD COLUMN "headers" text;`
        );

        const usedNiceIds: string[] = [];

        for (const resource of resources.rows) {
            // Generate a unique name and ensure it's unique
            let niceId = "";
            let loops = 0;
            while (true) {
                if (loops > 100) {
                    throw new Error("Could not generate a unique name");
                }

                niceId = generateName();
                if (!usedNiceIds.includes(niceId)) {
                    usedNiceIds.push(niceId);
                    break;
                }
                loops++;
            }
            await db.execute(sql`
                UPDATE "resources" SET "niceId" = ${niceId} WHERE "resourceId" = ${resource.resourceId}
            `);
        }

        for (const resource of siteResources.rows) {
            // Generate a unique name and ensure it's unique
            let niceId = "";
            let loops = 0;
            while (true) {
                if (loops > 100) {
                    throw new Error("Could not generate a unique name");
                }

                niceId = generateName();
                if (!usedNiceIds.includes(niceId)) {
                    usedNiceIds.push(niceId);
                    break;
                }
                loops++;
            }
            await db.execute(sql`
                UPDATE "siteResources" SET "niceId" = ${niceId} WHERE "siteResourceId" = ${resource.siteResourceId}
            `);
        }

        // Handle auto-provisioned users for identity providers
        const autoProvisionIdps = await db.execute(sql`
            SELECT "idpId" FROM "idp" WHERE "autoProvision" = true
        `);

        for (const idp of autoProvisionIdps.rows) {
            // Get all users with this identity provider
            const usersWithIdp = await db.execute(sql`
                SELECT "id" FROM "user" WHERE "idpId" = ${idp.idpId}
            `);

            // Update userOrgs to set autoProvisioned to true for these users
            for (const user of usersWithIdp.rows) {
                await db.execute(sql`
                    UPDATE "userOrgs" SET "autoProvisioned" = true WHERE "userId" = ${user.id}
                `);
            }
        }

        await db.execute(sql`COMMIT`);
        console.log(`Migrated database`);
    } catch (e) {
        await db.execute(sql`ROLLBACK`);
        console.log("Failed to migrate db:", e);
        throw e;
    }
}

const dev = process.env.ENVIRONMENT !== "prod";
let file;
if (!dev) {
    file = join(__DIRNAME, "names.json");
} else {
    file = join("server/db/names.json");
}
export const names = JSON.parse(readFileSync(file, "utf-8"));

export function generateName(): string {
    const name = (
        names.descriptors[
            Math.floor(Math.random() * names.descriptors.length)
        ] +
        "-" +
        names.animals[Math.floor(Math.random() * names.animals.length)]
    )
        .toLowerCase()
        .replace(/\s/g, "-");

    // clean out any non-alphanumeric characters except for dashes
    return name.replace(/[^a-z0-9-]/g, "");
}