From 1327b8251ba59356195a198ed4d91d362155d0f4 Mon Sep 17 00:00:00 2001
From: Thomas Bonnin <233326+TBonnin@users.noreply.github.com>
Date: Mon, 4 Mar 2024 12:14:06 +0100
Subject: [PATCH] retry db migrations

In case multiple instances of server are trying to run migration at
startup, knex lock mechanism can lead to one of the instance migration
to fail and therefore fail to deploy.
In order to avoid failed migration because of lock, we wait and retry
migrations
---
 packages/shared/lib/db/database.ts           |  6 ++-
 packages/shared/lib/utils/retry.ts           | 21 ++++++++++
 packages/shared/lib/utils/retry.unit.test.ts | 40 ++++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 packages/shared/lib/utils/retry.ts
 create mode 100644 packages/shared/lib/utils/retry.unit.test.ts

diff --git a/packages/shared/lib/db/database.ts b/packages/shared/lib/db/database.ts
index e771b8ff9c..db98566b3c 100644
--- a/packages/shared/lib/db/database.ts
+++ b/packages/shared/lib/db/database.ts
@@ -1,5 +1,6 @@
 import knex from 'knex';
 import type { Knex } from 'knex';
+import { retry } from '../utils/retry.js';
 
 export function getDbConfig({ timeoutMs }: { timeoutMs: number }): Knex.Config<any> {
     return {
@@ -30,7 +31,10 @@ export class KnexDatabase {
     }
 
     async migrate(directory: string): Promise<any> {
-        return this.knex.migrate.latest({ directory: directory, tableName: '_nango_auth_migrations', schemaName: this.schema() });
+        return retry(async () => await this.knex.migrate.latest({ directory: directory, tableName: '_nango_auth_migrations', schemaName: this.schema() }), {
+            maxAttempts: 4,
+            delayMs: (attempt) => 500 * attempt
+        });
     }
 
     schema() {
diff --git a/packages/shared/lib/utils/retry.ts b/packages/shared/lib/utils/retry.ts
new file mode 100644
index 0000000000..cf37224813
--- /dev/null
+++ b/packages/shared/lib/utils/retry.ts
@@ -0,0 +1,21 @@
+interface RetryConfig {
+    maxAttempts: number;
+    delayMs: number | ((attempt: number) => number);
+}
+
+export async function retry<T>(fn: () => T, config: RetryConfig): Promise<T> {
+    const { maxAttempts, delayMs } = config;
+    for (let attempt = 1; attempt <= maxAttempts; attempt++) {
+        try {
+            return fn();
+        } catch (error) {
+            if (attempt < maxAttempts) {
+                const delay = typeof delayMs === 'number' ? delayMs : delayMs(attempt);
+                await new Promise((resolve) => setTimeout(resolve, delay));
+            } else {
+                throw error;
+            }
+        }
+    }
+    throw new Error('unreachable');
+}
diff --git a/packages/shared/lib/utils/retry.unit.test.ts b/packages/shared/lib/utils/retry.unit.test.ts
new file mode 100644
index 0000000000..75a4e29d88
--- /dev/null
+++ b/packages/shared/lib/utils/retry.unit.test.ts
@@ -0,0 +1,40 @@
+import { expect, describe, it } from 'vitest';
+import { retry } from './retry';
+
+describe('retry', () => {
+    it('should retry', async () => {
+        let count = 0;
+        const result = await retry(
+            () => {
+                count++;
+                if (count < 3) {
+                    throw new Error('my error');
+                }
+                return count;
+            },
+            {
+                maxAttempts: 3,
+                delayMs: () => 0
+            }
+        );
+        expect(result).toEqual(3);
+    });
+
+    it('should throw error after max attempts', async () => {
+        let count = 0;
+        try {
+            await retry(
+                () => {
+                    count++;
+                    throw new Error('my error');
+                },
+                {
+                    maxAttempts: 3,
+                    delayMs: () => 0
+                }
+            );
+        } catch (error: any) {
+            expect(error.message).toEqual('my error');
+        }
+    });
+});