merged

1 files changed, 54 insertions, 2 deletions

diff --git a/src/server/database.ts b/src/server/database.ts
index 7f5331998..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,6 +150,34 @@ export class Database {
         }
     }
 
+    public async visit(ids: string[], fn: (result: any) => string[], collectionName = "newDocuments"): Promise<void> {
+        if (this.db) {
+            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 => Database.Instance.getDocuments(fetchIds, res, "newDocuments"));
+                for (const doc of docs) {
+                    const id = doc.id;
+                    visited.add(id);
+                    ids.push(...fn(doc));
+                }
+            }
+
+        } else {
+            return new Promise(res => {
+                this.onConnect.push(() => {
+                    this.visit(ids, fn, collectionName);
+                    res();
+                });
+            });
+        }
+    }
+
     public query(query: { [key: string]: any }, projection?: { [key: string]: 0 | 1 }, collectionName = "newDocuments"): Promise<mongodb.Cursor> {
         if (this.db) {
             let cursor = this.db.collection(collectionName).find(query);