3 files changed, 98 insertions, 3 deletions

diff --git a/src/server/IDatabase.ts b/src/server/IDatabase.ts
new file mode 100644
index 000000000..a7e3e2b0a
--- /dev/null
+++ b/src/server/IDatabase.ts
@@ -0,0 +1,21 @@
+import * as mongodb from 'mongodb';
+import { Transferable } from './Message';
+
+export const DocumentsCollection = 'documents';
+export const NewDocumentsCollection = 'newDocuments';
+export interface IDatabase {
+    update(id: string, value: any, callback: (err: mongodb.MongoError, res: mongodb.UpdateWriteOpResult) => void, upsert?: boolean, collectionName?: string): Promise<void>;
+
+    replace(id: string, value: any, callback: (err: mongodb.MongoError, res: mongodb.UpdateWriteOpResult) => void, upsert?: boolean, collectionName?: string): void;
+
+    delete(query: any, collectionName?: string): Promise<mongodb.DeleteWriteOpResultObject>;
+    delete(id: string, collectionName?: string): Promise<mongodb.DeleteWriteOpResultObject>;
+
+    deleteAll(collectionName?: string, persist?: boolean): Promise<any>;
+
+    insert(value: any, collectionName?: string): Promise<void>;
+
+    getDocument(id: string, fn: (result?: Transferable) => void, collectionName?: string): void;
+    getDocuments(ids: string[], fn: (result?: Transferable[]) => void, collectionName?: string): void;
+    visit(ids: string[], fn: (result: any) => string[] | Promise<string[]>, collectionName?: string): Promise<void>;
+}
diff --git a/src/server/MemoryDatabase.ts b/src/server/MemoryDatabase.ts
new file mode 100644
index 000000000..8bf6abdd8
--- /dev/null
+++ b/src/server/MemoryDatabase.ts
@@ -0,0 +1,73 @@
+import { IDatabase, DocumentsCollection, NewDocumentsCollection } from './IDatabase';
+import { Transferable } from './Message';
+import * as mongodb from 'mongodb';
+
+export class MemoryDatabase implements IDatabase {
+
+    private db: { [collectionName: string]: { [id: string]: any } } = {};
+
+    private getCollection(collectionName: string) {
+        let collection = this.db[collectionName];
+        if (collection) {
+            return collection;
+        } else {
+            return this.db[collectionName] = {};
+        }
+    }
+
+    public update(id: string, value: any, callback: (err: mongodb.MongoError, res: mongodb.UpdateWriteOpResult) => void, _upsert?: boolean, collectionName = DocumentsCollection): Promise<void> {
+        const collection = this.getCollection(collectionName);
+        collection[id] = { ...collection[id], ...value };
+        callback(null, {} as any);
+        return Promise.resolve(undefined);
+    }
+
+    public replace(id: string, value: any, callback: (err: mongodb.MongoError, res: mongodb.UpdateWriteOpResult) => void, upsert?: boolean, collectionName = DocumentsCollection): void {
+        this.update(id, value, callback, upsert, collectionName);
+    }
+
+    public delete(query: any, collectionName?: string): Promise<mongodb.DeleteWriteOpResultObject>;
+    public delete(id: string, collectionName?: string): Promise<mongodb.DeleteWriteOpResultObject>;
+    public delete(id: any, collectionName = DocumentsCollection): Promise<mongodb.DeleteWriteOpResultObject> {
+        const i = id["_id"] ?? id;
+        delete this.getCollection(collectionName)[i];
+
+        return Promise.resolve({} as any);
+    }
+
+    public deleteAll(collectionName = DocumentsCollection, _persist = true): Promise<any> {
+        delete this.db[collectionName];
+        return Promise.resolve();
+    }
+
+    public insert(value: any, collectionName = DocumentsCollection): Promise<void> {
+        const id = value._id;
+        this.getCollection(collectionName)[id] = value;
+        return Promise.resolve();
+    }
+
+    public getDocument(id: string, fn: (result?: Transferable) => void, collectionName = NewDocumentsCollection): void {
+        fn(this.getCollection(collectionName)[id]);
+    }
+    public getDocuments(ids: string[], fn: (result?: Transferable[]) => void, collectionName = DocumentsCollection): void {
+        fn(ids.map(id => this.getCollection(collectionName)[id]));
+    }
+
+    public async visit(ids: string[], fn: (result: any) => string[] | Promise<string[]>, collectionName = NewDocumentsCollection): Promise<void> {
+        const visited = new Set<string>();
+        while (ids.length) {
+            const count = Math.min(ids.length, 1000);
+            const index = ids.length - count;
+            const fetchIds = ids.splice(index, count).filter(id => !visited.has(id));
+            if (!fetchIds.length) {
+                continue;
+            }
+            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);
+                ids.push(...(await fn(doc)));
+            }
+        }
+    }
+}
diff --git a/src/server/database.ts b/src/server/database.ts
index 6e0771c11..3039baedc 100644
--- a/src/server/database.ts
+++ b/src/server/database.ts
@@ -5,6 +5,7 @@ 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 * as mongoose from 'mongoose';
 
 export namespace Database {
@@ -44,7 +45,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 +216,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);
@@ -331,4 +332,4 @@ export namespace Database {
 
     }
 
-}
\ No newline at end of file
+}