1 files changed, 17 insertions, 6 deletions

diff --git a/src/server/database.ts b/src/server/database.ts
index 5bdf1fc45..83ce865c6 100644
--- a/src/server/database.ts
+++ b/src/server/database.ts
@@ -5,11 +5,13 @@ import { Utils, emptyFunction } from '../Utils';
 import { DashUploadUtils } from './DashUploadUtils';
 import { Credentials } from 'google-auth-library';
 import { GoogleApiServerUtils } from './apis/google/GoogleApiServerUtils';
+import { IDatabase } from './IDatabase';
+import { MemoryDatabase } from './MemoryDatabase';
 import * as mongoose from 'mongoose';
-import { addBeforeExitHandler } from './ActionUtilities';
 
 export namespace Database {
 
+    export let disconnect: Function;
     const schema = 'Dash';
     const port = 27017;
     export const url = `mongodb://localhost:${port}/${schema}`;
@@ -25,7 +27,7 @@ export namespace Database {
     export async function tryInitializeConnection() {
         try {
             const { connection } = mongoose;
-            addBeforeExitHandler(async () => { await new Promise<any>(resolve => connection.close(resolve)); });
+            disconnect = async () => new Promise<any>(resolve => connection.close(resolve));
             if (connection.readyState === ConnectionStates.disconnected) {
                 await new Promise<void>((resolve, reject) => {
                     connection.on('error', reject);
@@ -44,7 +46,7 @@ export namespace Database {
         }
     }
 
-    class Database {
+    class Database implements IDatabase {
         public static DocumentsCollection = 'documents';
         private MongoClient = mongodb.MongoClient;
         private currentWrites: { [id: string]: Promise<void> } = {};
@@ -215,7 +217,7 @@ export namespace Database {
                     if (!fetchIds.length) {
                         continue;
                     }
-                    const docs = await new Promise<{ [key: string]: any }[]>(res => Instance.getDocuments(fetchIds, res, "newDocuments"));
+                    const docs = await new Promise<{ [key: string]: any }[]>(res => this.getDocuments(fetchIds, res, collectionName));
                     for (const doc of docs) {
                         const id = doc.id;
                         visited.add(id);
@@ -262,7 +264,16 @@ export namespace Database {
         }
     }
 
-    export const Instance = new Database();
+    function getDatabase() {
+        switch (process.env.DB) {
+            case "MEM":
+                return new MemoryDatabase();
+            default:
+                return new Database();
+        }
+    }
+
+    export const Instance: IDatabase = getDatabase();
 
     export namespace Auxiliary {
 
@@ -331,4 +342,4 @@ export namespace Database {
 
     }
 
-}
\ No newline at end of file
+}