From 8fa3ec895fd6d6d6c8bb5eca73b9f4a87791caaf Mon Sep 17 00:00:00 2001
From: Eric Peterson <ericpeterson@spotify.com>
Date: Tue, 1 Oct 2024 14:23:22 +0200
Subject: [PATCH] Refactor DB Manager shutdown to be private

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

diff --git a/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts b/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts
index ff95c0562a..9441c8bca9 100644
--- a/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts
+++ b/packages/backend-defaults/src/entrypoints/database/DatabaseManager.ts
@@ -19,6 +19,8 @@ import {
   LifecycleService,
   LoggerService,
   RootConfigService,
+  RootLifecycleService,
+  RootLoggerService,
 } from '@backstage/backend-plugin-api';
 import { Config } from '@backstage/config';
 import { stringifyError } from '@backstage/errors';
@@ -42,6 +44,8 @@ function pluginPath(pluginId: string): string {
  */
 export type DatabaseManagerOptions = {
   migrations?: DatabaseService['migrations'];
+  rootLogger?: RootLoggerService;
+  rootLifecycle?: RootLifecycleService;
 };
 
 /**
@@ -53,7 +57,15 @@ export class DatabaseManagerImpl {
     private readonly connectors: Record<string, Connector>,
     private readonly options?: DatabaseManagerOptions,
     private readonly databaseCache: Map<string, Promise<Knex>> = new Map(),
-  ) {}
+  ) {
+    // If a rootLifecycle service was provided, register a shutdown hook to
+    // clean up any database connections.
+    if (options?.rootLifecycle !== undefined) {
+      options.rootLifecycle.addShutdownHook(async () => {
+        await this.shutdown({ logger: options.rootLogger });
+      });
+    }
+  }
 
   /**
    * Generates a PluginDatabaseManager for consumption by plugins.
@@ -90,16 +102,16 @@ export class DatabaseManagerImpl {
   }
 
   /**
-   * Method to be called during shutdown to destroy all known connections.
+   * Destroys all known connections.
    */
-  async shutdown(deps: { logger: LoggerService }): Promise<void> {
+  private async shutdown(deps?: { logger?: LoggerService }): Promise<void> {
     const pluginIds = Array.from(this.databaseCache.keys());
     await Promise.allSettled(
       pluginIds.map(async pluginId => {
         const connection = await this.databaseCache.get(pluginId);
         if (connection) {
           await connection.destroy().catch((error: unknown) => {
-            deps.logger.error(
+            deps?.logger?.error(
               `Problem closing database connection for ${pluginId}: ${stringifyError(
                 error,
               )}`,
@@ -257,11 +269,4 @@ 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 5e23e6ac4c..525830471b 100644
--- a/packages/backend-defaults/src/entrypoints/database/databaseServiceFactory.ts
+++ b/packages/backend-defaults/src/entrypoints/database/databaseServiceFactory.ts
@@ -41,21 +41,16 @@ export const databaseServiceFactory = createServiceFactory({
     rootLogger: coreServices.rootLogger,
   },
   async createRootContext({ config, rootLifecycle, rootLogger }) {
-    const databaseManager = config.getOptional('backend.database')
-      ? DatabaseManager.fromConfig(config)
+    return config.getOptional('backend.database')
+      ? DatabaseManager.fromConfig(config, { rootLifecycle, rootLogger })
       : DatabaseManager.fromConfig(
           new ConfigReader({
             backend: {
               database: { client: 'better-sqlite3', connection: ':memory:' },
             },
           }),
+          { rootLifecycle, rootLogger },
         );
-
-    rootLifecycle.addShutdownHook(async () => {
-      await databaseManager.shutdown({ logger: rootLogger });
-    });
-
-    return databaseManager;
   },
   async factory({ pluginMetadata, lifecycle, logger }, databaseManager) {
     return databaseManager.forPlugin(pluginMetadata.getId(), {