From 860ebb73a9552cff980a9d09ff5cfba20288cc91 Mon Sep 17 00:00:00 2001
From: Daniel Hougaard <daniel@infisical.com>
Date: Sat, 2 Nov 2024 19:44:09 +0400
Subject: [PATCH] Update 20241014084900_identity-multiple-auth-methods.ts

---
 ...14084900_identity-multiple-auth-methods.ts | 34 ++++++-------------
 1 file changed, 11 insertions(+), 23 deletions(-)

diff --git a/backend/src/db/migrations/20241014084900_identity-multiple-auth-methods.ts b/backend/src/db/migrations/20241014084900_identity-multiple-auth-methods.ts
index b669cff53..35d1984de 100644
--- a/backend/src/db/migrations/20241014084900_identity-multiple-auth-methods.ts
+++ b/backend/src/db/migrations/20241014084900_identity-multiple-auth-methods.ts
@@ -2,7 +2,7 @@ import { Knex } from "knex";
 
 import { TableName } from "../schemas";
 
-const BATCH_SIZE = 10000;
+const BATCH_SIZE = 10_000;
 
 export async function up(knex: Knex): Promise<void> {
   const hasAuthMethodColumnAccessToken = await knex.schema.hasColumn(TableName.IdentityAccessToken, "authMethod");
@@ -12,25 +12,11 @@ export async function up(knex: Knex): Promise<void> {
       t.string("authMethod").nullable();
     });
 
-    // Get total count of records to process
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore because count is a string
-    const count = (await knex(TableName.IdentityAccessToken).count("id as count").first()) as { count: string };
+    let nullableAccessTokens = await knex(TableName.IdentityAccessToken).whereNull("authMethod").limit(BATCH_SIZE);
+    let totalUpdated = 0;
 
-    if (!count) {
-      throw new Error("Failed to find count of identity access tokens");
-    }
-
-    const totalRecords = parseInt(count.count, 10);
-    // Process in batches
-    for (let offset = 0; offset < totalRecords; offset += BATCH_SIZE) {
-      // ! Get the current access tokens to process
-      // eslint-disable-next-line no-await-in-loop
-      const batchIds = await knex(TableName.IdentityAccessToken)
-        .select("id")
-        .limit(BATCH_SIZE)
-        .offset(offset)
-        .pluck("id");
+    do {
+      const batchIds = nullableAccessTokens.map((token) => token.id);
 
       // ! Update the auth method column in batches for the current batch
       // eslint-disable-next-line no-await-in-loop
@@ -46,10 +32,12 @@ export async function up(knex: Knex): Promise<void> {
             .first()
         });
 
-      // Log progress
-      // eslint-disable-next-line no-console
-      console.log(`Processed ${Math.min(offset + BATCH_SIZE, totalRecords)} of ${totalRecords} records`);
-    }
+      // eslint-disable-next-line no-await-in-loop
+      nullableAccessTokens = await knex(TableName.IdentityAccessToken).whereNull("authMethod").limit(BATCH_SIZE);
+
+      totalUpdated += batchIds.length;
+      console.log(`Updated ${batchIds.length} access tokens in batch <> Total updated: ${totalUpdated}`);
+    } while (nullableAccessTokens.length > 0);
 
     // ! We delete all access tokens where the identity has no auth method set!
     // ! Which means un-configured identities that for some reason have access tokens, will have their access tokens deleted.