core: db upd

packages/hydrooj/src/entry/worker.ts

@@ -25,6 +25,7 @@ export async function load() {
     const opts = options();
     await db.start(opts);
     const modelSystem = require('../model/system');
+    await modelSystem.runConfig();
     const [endPoint, accessKey, secretKey, bucket, region, endPointForUser, endPointForJudge] = modelSystem.getMany([
         'file.endPoint', 'file.accessKey', 'file.secretKey', 'file.bucket', 'file.region',
         'file.endPointForUser', 'file.endPointForJudge',

packages/hydrooj/src/model/system.ts

@@ -7,7 +7,9 @@ import * as bus from '../service/bus';
 const coll = db.collection('system');
 const cache: Record<string, any> = {};
 
-export function get<K extends keyof SystemKeys>(key: K): SystemKeys[K] {
+export function get<K extends keyof SystemKeys>(key: K): SystemKeys[K]
+export function get(key: string): any
+export function get(key: string): any {
     return cache[key];
 }
 
@@ -34,7 +36,9 @@ export function getMany(keys: string[]): any[] {
     return keys.map((key) => cache[key]);
 }
 
-export async function set<K extends keyof SystemKeys>(_id: K, value: SystemKeys[K]) {
+export async function set<K extends keyof SystemKeys>(_id: K, value: SystemKeys[K]): Promise<SystemKeys[K]>
+export async function set<K>(_id: string, value: K): Promise<K>
+export async function set(_id: string, value: any) {
     const res = await coll.findOneAndUpdate(
         { _id },
         { $set: { value } },

packages/hydrooj/src/service/storage.ts

@@ -90,7 +90,8 @@ class StorageService {
         return await this.client.putObject(this.opts.bucket, target, file, meta);
     }
 
-    async get(target: string) {
+    async get(target: string, path?: string) {
+        if (path) return await this.client.fGetObject(this.opts.bucket, target, path);
         return await this.client.getObject(this.opts.bucket, target);
     }
 

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) {
+        terminal.clear('Changing storage engine. This may take a long time.');
-            const stream = gridfs.openDownloadStream(_id);
-            stream.on('error', (err) => {
-                console.error(err);
-                if (reject) reject(err);
-            });
-            return stream;
-        }
-        terminal.clear('Upgrade');
-        logger.info('');
         // Problem file and User file
-        // @ts-ignore
+        try {
-        const progress = terminal.progressBar({ title: 'Changing storage engine. This may take a long time.', y: 1 });
+            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({
-            // @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();
+                const pdocs = await problem.getMulti(ddoc._id, { data: { $ne: null } }, ['domainId', 'docId', 'data', 'title']).toArray();
-            // @ts-ignore
                 const domainProgress = terminal.progressBar({
-                // @ts-ignore
                     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}`);
-                // TODO copy file
+                    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);
             }
             totalProgress.stop();
-        // TODO userfile
+            const udocs = await user.getMulti().project({ _id: 1, uname: 1 }).toArray();
-        progress.stop();
+            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;
+        }
     },
 ];