From 441af526e4e6d246615d39586ad939f91de239ff Mon Sep 17 00:00:00 2001
From: Thomas Bonnin <233326+TBonnin@users.noreply.github.com>
Date: Mon, 4 Mar 2024 14:00:54 +0100
Subject: [PATCH] retry db migrations (#1791)

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

## Checklist before requesting a review (skip if just adding/editing
APIs & templates)
- [x] I added tests, otherwise the reason is:
- [ ] I added observability, otherwise the reason is:
- [ ] I added analytics, otherwise the reason is:
---
 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');
+        }
+    });
+});