1 files changed, 33 insertions, 4 deletions
diff --git a/src/server/database.ts b/src/server/database.ts
index acb6ce751..a7254fb0c 100644
--- a/src/server/database.ts
+++ b/src/server/database.ts
@@ -17,7 +17,7 @@ export class Database {
         });
     }
 
-    public update(id: string, value: any, callback: () => void, upsert = true, collectionName = Database.DocumentsCollection) {
+    public update(id: string, value: any, callback: (err: mongodb.MongoError, res: mongodb.UpdateWriteOpResult) => void, upsert = true, collectionName = Database.DocumentsCollection) {
         if (this.db) {
             let collection = this.db.collection(collectionName);
             const prom = this.currentWrites[id];
@@ -30,7 +30,7 @@ export class Database {
                                 delete this.currentWrites[id];
                             }
                             resolve();
-                            callback();
+                            callback(err, res);
                         });
                 });
             };
@@ -41,6 +41,30 @@ export class Database {
         }
     }
 
+    public replace(id: string, value: any, callback: (err: mongodb.MongoError, res: mongodb.UpdateWriteOpResult) => void, upsert = true, collectionName = Database.DocumentsCollection) {
+        if (this.db) {
+            let collection = this.db.collection(collectionName);
+            const prom = this.currentWrites[id];
+            let newProm: Promise<void>;
+            const run = (): Promise<void> => {
+                return new Promise<void>(resolve => {
+                    collection.replaceOne({ _id: id }, value, { upsert }
+                        , (err, res) => {
+                            if (this.currentWrites[id] === newProm) {
+                                delete this.currentWrites[id];
+                            }
+                            resolve();
+                            callback(err, res);
+                        });
+                });
+            };
+            newProm = prom ? prom.then(run) : run();
+            this.currentWrites[id] = newProm;
+        } else {
+            this.onConnect.push(() => this.replace(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 = Database.DocumentsCollection) {
@@ -126,7 +150,7 @@ export class Database {
         }
     }
 
-    public async visit(ids: string[], fn: (result: any) => string[], collectionName = "newDocuments") {
+    public async visit(ids: string[], fn: (result: any) => string[], collectionName = "newDocuments"): Promise<void> {
         if (this.db) {
             const visited = new Set<string>();
             while (ids.length) {
@@ -145,7 +169,12 @@ export class Database {
             }
 
         } else {
-            this.onConnect.push(() => this.visit(ids, fn, collectionName));
+            return new Promise(res => {
+                this.onConnect.push(() => {
+                    this.visit(ids, fn, collectionName);
+                    res();
+                });
+            });
         }
     }