core: db upd

packages/hydrooj/src/upgrade.ts

@@ -1,14 +1,24 @@
 /* eslint-disable no-await-in-loop */
 import { ObjectID, Collection } from 'mongodb';
 import { terminal } from 'terminal-kit';
+import AdmZip from 'adm-zip';
 import { Ufdoc } from './interface';
 import { Logger } from './logger';
+import { streamToBuffer } from './utils';
 import gridfs from './service/gridfs';
 import storage from './service/storage';
 import db from './service/db';
 import * as problem from './model/problem';
+import * as user from './model/user';
 import * as domain from './model/domain';
 import * as document from './model/document';
+import * as system from './model/system';
+
+declare module 'terminal-kit/Terminal' {
+    interface ProgressBarOptions {
+        y?: number;
+    }
+}
 
 const logger = new Logger('upgrade');
 type UpgradeScript = () => Promise<boolean | void>;
@@ -49,7 +59,6 @@ const scripts: UpgradeScript[] = [
             logger.error('Cannot upgrade. Please change storage config.');
             return false;
         }
-
         const coll: Collection<Ufdoc> = db.collection('file');
         const collFile = db.collection('fs.files');
         const collChunk = db.collection('fs.chunks');
@@ -60,46 +69,80 @@ const scripts: UpgradeScript[] = [
                 collChunk.deleteMany({ files_id: _id }),
             ]);
         }
-        async function get(_id: ObjectID, reject?: Function) {
-            const stream = gridfs.openDownloadStream(_id);
-            stream.on('error', (err) => {
-                console.error(err);
-                if (reject) reject(err);
-            });
-            return stream;
-        }
-        terminal.clear('Upgrade');
-        logger.info('');
+        terminal.clear('Changing storage engine. This may take a long time.');
         // Problem file and User file
-        // @ts-ignore
-        const progress = terminal.progressBar({ title: 'Changing storage engine. This may take a long time.', y: 1 });
-        const ddocs = await domain.getMulti().toArray();
-        const totalProgress = terminal.progressBar({
-            // @ts-ignore
-            items: ddocs.length, title: 'Domain', y: 3, eta: true,
-        });
-        for (let i = 0; i < ddocs.length; i++) {
-            const ddoc = ddocs[i];
-            totalProgress.startItem(ddoc._id);
-            const pdocs = await problem.getMulti(ddoc._id, { data: { $ne: null } }, ['docId', 'data', 'title']).toArray();
-            // @ts-ignore
-            const domainProgress = terminal.progressBar({
-                // @ts-ignore
-                items: pdocs.length, title: 'Problems', y: 4, eta: true,
+        try {
+            let savedProgress = system.get('upgrade.file.progress.domain');
+            if (savedProgress) savedProgress = JSON.parse(savedProgress);
+            else savedProgress = { pdocs: [] };
+            const ddocs = await domain.getMulti().toArray();
+            const totalProgress = terminal.progressBar({
+                items: ddocs.length, title: 'Domain', y: 3, eta: true,
             });
-            for (let j = 0; j < pdocs.length; j++) {
-                const pdoc = pdocs[j];
-                domainProgress.startItem(`${pdoc.docId}: ${pdoc.title}`);
-                // TODO copy file
-                domainProgress.itemDone(`${pdoc.docId}: ${pdoc.title}`);
+            for (let i = 0; i < ddocs.length; i++) {
+                const ddoc = ddocs[i];
+                totalProgress.startItem(ddoc._id);
+                const pdocs = await problem.getMulti(ddoc._id, { data: { $ne: null } }, ['domainId', 'docId', 'data', 'title']).toArray();
+                const domainProgress = terminal.progressBar({
+                    items: pdocs.length, title: 'Problems', y: 4, eta: true,
+                });
+                for (let j = 0; j < pdocs.length; j++) {
+                    const pdoc = pdocs[j];
+                    domainProgress.startItem(`${pdoc.docId}: ${pdoc.title}`);
+                    if (!savedProgress.pdocs.includes(`${pdoc.domainId}/${pdoc.docId}`) && pdoc.data instanceof ObjectID) {
+                        try {
+                            const file = await streamToBuffer(gridfs.openDownloadStream(pdoc.data));
+                            const zip = new AdmZip(file);
+                            const entries = zip.getEntries();
+                            await Promise.all(entries.map(
+                                (entry) => storage.put(`problem/${pdoc.domainId}/${pdoc.docId}/testdata/${entry.name}`, entry.getData()),
+                            ));
+                        } catch (e) {
+                            if (e.toString().includes('FileNotFound')) {
+                                logger.error('\n\n\n\n\nProblem data not found %s/%s', pdoc.domainId, pdoc.docId);
+                            } else throw e;
+                        }
+                        savedProgress.pdocs.push(`${pdoc.domainId}/${pdoc.docId}`);
+                    }
+                    system.set('upgrade.file.progress.domain', JSON.stringify(savedProgress));
+                    domainProgress.itemDone(`${pdoc.docId}: ${pdoc.title}`);
+                }
+                domainProgress.stop();
+                totalProgress.itemDone(ddoc._id);
             }
-            domainProgress.stop();
-            totalProgress.itemDone(ddoc._id);
+            totalProgress.stop();
+            const udocs = await user.getMulti().project({ _id: 1, uname: 1 }).toArray();
+            const userfileProgress = terminal.progressBar({
+                items: udocs.length, title: 'Users', y: 3,
+            });
+            for (const udoc of udocs) {
+                userfileProgress.startItem(`${udoc._id}: ${udoc.uname}`);
+                const ufdocs = await coll.find({ owner: udoc._id }).toArray();
+                const currentUser = terminal.progressBar({
+                    items: ufdocs.length, title: `${udoc._id}: ${udoc.uname}`, y: 4, eta: true,
+                });
+                const namelist = [];
+                for (const ufdoc of ufdocs) {
+                    currentUser.startItem(`${ufdoc._id}: ${ufdoc.filename || ''}`);
+                    ufdoc.filename = ufdoc.filename || ufdoc._id.toHexString();
+                    const file = await streamToBuffer(gridfs.openDownloadStream(ufdoc._id));
+                    if (!namelist.includes(ufdoc.filename)) {
+                        await storage.put(`user/${udoc._id}/${ufdoc.filename}`, file);
+                    } else {
+                        await storage.put(`user/${udoc._id}/${ufdoc._id}`, file);
+                    }
+                    namelist.push(ufdoc.filename);
+                    currentUser.itemDone(`${ufdoc._id}: ${ufdoc.filename || ''}`);
+                }
+                currentUser.stop();
+                userfileProgress.itemDone(`${udoc._id}: ${udoc.uname}`);
+            }
+            userfileProgress.stop();
+            return true;
+        } catch (e) {
+            logger.error('\n\n\n\n\n%o', e);
+            return false;
         }
-        totalProgress.stop();
-        // TODO userfile
-        progress.stop();
-        return false;
     },
 ];