From 35db939199a2d826e7ee4dbe31af48cc42364ea2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20C=C3=A1rdenas?= <rafael@rafaelcr.com>
Date: Wed, 21 Sep 2022 10:50:23 -0500
Subject: [PATCH] fix: retry pg connection on new library code (#1326)

* fix: retry pg connection on new library code

* fix: add ECONNRESET error code
---
 src/datastore/connection-legacy.ts | 31 +----------------------------
 src/datastore/connection.ts        |  5 +++--
 src/datastore/helpers.ts           | 32 ++++++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 32 deletions(-)

diff --git a/src/datastore/connection-legacy.ts b/src/datastore/connection-legacy.ts
index 8327667f..9a63552f 100644
--- a/src/datastore/connection-legacy.ts
+++ b/src/datastore/connection-legacy.ts
@@ -1,6 +1,7 @@
 import { Client, ClientConfig, Pool, PoolClient, PoolConfig } from 'pg';
 import { logError, logger, parseArgBoolean, parsePort, stopwatch, timeout } from '../helpers';
 import { PgServer } from './connection';
+import { isPgConnectionError } from './helpers';
 
 export type PgClientConfig = ClientConfig & { schema?: string };
 type PgPoolConfig = PoolConfig & { schema?: string };
@@ -167,33 +168,3 @@ export async function connectWithRetry(pool: Pool): Promise<PoolClient> {
     }
   }
 }
-
-/**
- * Checks if a given error from the pg lib is a connection error (i.e. the query is retryable).
- * If true then returns a normalized error message, otherwise returns false.
- */
-function isPgConnectionError(error: any): string | false {
-  if (error.code === 'ECONNREFUSED') {
-    return 'Postgres connection ECONNREFUSED';
-  } else if (error.code === 'ETIMEDOUT') {
-    return 'Postgres connection ETIMEDOUT';
-  } else if (error.code === 'ENOTFOUND') {
-    return 'Postgres connection ENOTFOUND';
-  } else if (error.message) {
-    const msg = (error as Error).message.toLowerCase();
-    if (msg.includes('database system is starting up')) {
-      return 'Postgres connection failed while database system is starting up';
-    } else if (msg.includes('database system is shutting down')) {
-      return 'Postgres connection failed while database system is shutting down';
-    } else if (msg.includes('connection terminated unexpectedly')) {
-      return 'Postgres connection terminated unexpectedly';
-    } else if (msg.includes('connection terminated')) {
-      return 'Postgres connection terminated';
-    } else if (msg.includes('connection error')) {
-      return 'Postgres client has encountered a connection error and is not queryable';
-    } else if (msg.includes('terminating connection due to unexpected postmaster exit')) {
-      return 'Postgres connection terminating due to unexpected postmaster exit';
-    }
-  }
-  return false;
-}
diff --git a/src/datastore/connection.ts b/src/datastore/connection.ts
index 0c5c903d..0837dc52 100644
--- a/src/datastore/connection.ts
+++ b/src/datastore/connection.ts
@@ -1,5 +1,6 @@
-import { has0xPrefix, logError, parseArgBoolean, parsePort, stopwatch, timeout } from '../helpers';
+import { logError, parseArgBoolean, parsePort, stopwatch, timeout } from '../helpers';
 import * as postgres from 'postgres';
+import { isPgConnectionError } from './helpers';
 
 export type PgSqlClient = postgres.Sql<any>;
 
@@ -87,7 +88,7 @@ export async function connectPostgres({
       connectionOkay = true;
       break;
     } catch (error: any) {
-      if (error instanceof postgres.PostgresError) {
+      if (isPgConnectionError(error) || error instanceof postgres.PostgresError) {
         const timeElapsed = initTimer.getElapsed();
         if (timeElapsed - lastElapsedLog > 2000) {
           lastElapsedLog = timeElapsed;
diff --git a/src/datastore/helpers.ts b/src/datastore/helpers.ts
index 6358a083..b8612676 100644
--- a/src/datastore/helpers.ts
+++ b/src/datastore/helpers.ts
@@ -171,6 +171,38 @@ export const TX_METADATA_TABLES = [
   'subdomains',
 ] as const;
 
+/**
+ * Checks if a given error from the pg lib is a connection error (i.e. the query is retryable).
+ * If true then returns a normalized error message, otherwise returns false.
+ */
+export function isPgConnectionError(error: any): string | false {
+  if (error.code === 'ECONNREFUSED') {
+    return 'Postgres connection ECONNREFUSED';
+  } else if (error.code === 'ETIMEDOUT') {
+    return 'Postgres connection ETIMEDOUT';
+  } else if (error.code === 'ENOTFOUND') {
+    return 'Postgres connection ENOTFOUND';
+  } else if (error.code === 'ECONNRESET') {
+    return 'Postgres connection ECONNRESET';
+  } else if (error.message) {
+    const msg = (error as Error).message.toLowerCase();
+    if (msg.includes('database system is starting up')) {
+      return 'Postgres connection failed while database system is starting up';
+    } else if (msg.includes('database system is shutting down')) {
+      return 'Postgres connection failed while database system is shutting down';
+    } else if (msg.includes('connection terminated unexpectedly')) {
+      return 'Postgres connection terminated unexpectedly';
+    } else if (msg.includes('connection terminated')) {
+      return 'Postgres connection terminated';
+    } else if (msg.includes('connection error')) {
+      return 'Postgres client has encountered a connection error and is not queryable';
+    } else if (msg.includes('terminating connection due to unexpected postmaster exit')) {
+      return 'Postgres connection terminating due to unexpected postmaster exit';
+    }
+  }
+  return false;
+}
+
 /**
  * Adds a table name prefix to an array of column names.
  * @param columns - array of column names