From b1254b40a8b78782acc4ecace2e118b8362f3b75 Mon Sep 17 00:00:00 2001
From: Eric Peterson <ericpeterson@spotify.com>
Date: Tue, 1 Oct 2024 12:57:15 +0200
Subject: [PATCH] Add root lifecycle shutdown hook to clean up DB connections

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
---
 .../entrypoints/database/DatabaseManager.ts   | 28 +++++++++++++++++++
 .../database/databaseServiceFactory.ts        | 12 ++++++--
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts b/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts
index 46d4ee8cf7..6ed5545d35 100644
--- a/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts
+++ b/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts
@@ -89,6 +89,27 @@ export class DatabaseManagerImpl {
     return { getClient, migrations: { skip } };
   }
 
+  /**
+   * Method to be called during shutdown to destroy all known connections.
+   */
+  async shutdown(deps: { logger: LoggerService }): Promise<void> {
+    const pluginIds = Array.from(this.databaseCache.keys());
+    await Promise.all(
+      pluginIds.map(async pluginId => {
+        const connection = await this.databaseCache.get(pluginId);
+        if (connection) {
+          await connection.destroy().catch((error: unknown) => {
+            deps.logger.error(
+              `Problem closing database connection for ${pluginId}: ${stringifyError(
+                error,
+              )}`,
+            );
+          });
+        }
+      }),
+    );
+  }
+
   /**
    * Provides the client type which should be used for a given plugin.
    *
@@ -236,4 +257,11 @@ export class DatabaseManager {
   ): PluginDatabaseManager {
     return this.impl.forPlugin(pluginId, deps);
   }
+
+  /**
+   * Method to be called during shutdown to destroy all known connections.
+   */
+  async shutdown(deps: { logger: LoggerService }): Promise<void> {
+    return this.impl.shutdown(deps);
+  }
 }
diff --git a/packages/backend-defaults/src/entrypoints/database/databaseServiceFactory.ts b/packages/backend-defaults/src/entrypoints/database/databaseServiceFactory.ts
index 167779e807..5e23e6ac4c 100644
--- a/packages/backend-defaults/src/entrypoints/database/databaseServiceFactory.ts
+++ b/packages/backend-defaults/src/entrypoints/database/databaseServiceFactory.ts
@@ -37,9 +37,11 @@ export const databaseServiceFactory = createServiceFactory({
     lifecycle: coreServices.lifecycle,
     logger: coreServices.logger,
     pluginMetadata: coreServices.pluginMetadata,
+    rootLifecycle: coreServices.rootLifecycle,
+    rootLogger: coreServices.rootLogger,
   },
-  async createRootContext({ config }) {
-    return config.getOptional('backend.database')
+  async createRootContext({ config, rootLifecycle, rootLogger }) {
+    const databaseManager = config.getOptional('backend.database')
       ? DatabaseManager.fromConfig(config)
       : DatabaseManager.fromConfig(
           new ConfigReader({
@@ -48,6 +50,12 @@ export const databaseServiceFactory = createServiceFactory({
             },
           }),
         );
+
+    rootLifecycle.addShutdownHook(async () => {
+      await databaseManager.shutdown({ logger: rootLogger });
+    });
+
+    return databaseManager;
   },
   async factory({ pluginMetadata, lifecycle, logger }, databaseManager) {
     return databaseManager.forPlugin(pluginMetadata.getId(), {