From d33a3a1fb0bb6b8738555f05a7fc506eb8466f14 Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Mon, 12 Aug 2024 03:36:41 +0800
Subject: [PATCH] core: use problem import options

---
 packages/hydrooj/src/handler/import.ts |  4 +--
 packages/hydrooj/src/model/problem.ts  | 47 ++++++++++++++++++++------
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/packages/hydrooj/src/handler/import.ts b/packages/hydrooj/src/handler/import.ts
index 844ed0eb4..e1f7faa4b 100644
--- a/packages/hydrooj/src/handler/import.ts
+++ b/packages/hydrooj/src/handler/import.ts
@@ -18,8 +18,8 @@ class ProblemImportHydroHandler extends Handler {
         if (!this.request.files.file) throw new ValidationError('file');
         if (preferredPrefix && !/^[a-zA-Z]+$/.test(preferredPrefix)) throw new ValidationError('preferredPrefix');
         const promise = problem.import(
-            domainId, this.request.files.file.filepath, keepUser ? null : this.user._id,
-            preferredPrefix, this.progress.bind(this),
+            domainId, this.request.files.file.filepath,
+            { preferredPrefix, progress: this.progress.bind(this), operator: keepUser ? null : this.user._id },
         ).catch((e) => MessageModel.send(1, this.user._id, `Import failed: ${e.message}\n${e.stack}`));
         let resolved = false;
         await Promise.race([
diff --git a/packages/hydrooj/src/model/problem.ts b/packages/hydrooj/src/model/problem.ts
index 6b89fb096..c788af107 100644
--- a/packages/hydrooj/src/model/problem.ts
+++ b/packages/hydrooj/src/model/problem.ts
@@ -54,6 +54,13 @@ function findOverrideContent(dir: string, base: string) {
     return JSON.stringify(languages);
 }
 
+interface ProblemImportOptions {
+    preferredPrefix?: string;
+    progress?: any;
+    override?: boolean;
+    operator?: number;
+}
+
 export class ProblemModel {
     static PROJECTION_CONTEST_LIST: Field[] = [
         '_id', 'domainId', 'docType', 'docId', 'pid',
@@ -439,9 +446,16 @@ export class ProblemModel {
         return true;
     }
 
-    static async import(domainId: string, filepath: string, operator = 1, preferredPrefix?: string, progress = null) {
+    static async import(domainId: string, filepath: string, options: ProblemImportOptions = {}) {
         let tmpdir = '';
         let del = false;
+        if (typeof options !== 'object') {
+            logger.warn('ProblemModel.import: options should be an object');
+            options = {};
+        }
+        const {
+            preferredPrefix, progress, override = false, operator = 1,
+        } = options;
         try {
             if (filepath.endsWith('.zip')) {
                 tmpdir = path.join(os.tmpdir(), 'hydro', `${Math.random()}.import`);
@@ -474,10 +488,16 @@ export class ProblemModel {
                 const pdoc: ProblemDoc = yaml.load(content) as any;
                 if (!pdoc) continue;
                 let pid = pdoc.pid;
+                let overridePid = null;
 
                 const isValidPid = async (id: string) => {
                     if (!(/^[A-Za-z]+[0-9A-Za-z]*$/.test(id))) return false;
-                    if (await ProblemModel.get(domainId, id)) return false;
+                    const doc = await ProblemModel.get(domainId, id);
+                    if (doc) {
+                        if (!override) return false;
+                        overridePid = doc.docId;
+                        return true;
+                    }
                     return true;
                 };
                 const getFiles = async (type: string) => {
@@ -496,10 +516,18 @@ export class ProblemModel {
                 const overrideContent = findOverrideContent(path.join(tmpdir, i), 'problem');
                 if (pdoc.difficulty && !Number.isSafeInteger(pdoc.difficulty)) delete pdoc.difficulty;
                 if (typeof pdoc.title !== 'string') throw new ValidationError('title', null, 'Invalid title');
-                const docId = await ProblemModel.add(
-                    domainId, pid, pdoc.title.trim(), overrideContent || pdoc.content || 'No content',
-                    operator || pdoc.owner, pdoc.tag || [], { hidden: pdoc.hidden, difficulty: pdoc.difficulty },
-                );
+                const docId = overridePid
+                    ? (await ProblemModel.edit(domainId, overridePid, {
+                        title: pdoc.title.trim(),
+                        content: overrideContent || pdoc.content || 'No content',
+                        tag: pdoc.tag || [],
+                        difficulty: pdoc.difficulty,
+                    })).docId
+                    : await ProblemModel.add(
+                        domainId, pid, pdoc.title.trim(), overrideContent || pdoc.content || 'No content',
+                        operator || pdoc.owner, pdoc.tag || [], { hidden: pdoc.hidden, difficulty: pdoc.difficulty },
+                    );
+                // TODO delete unused file when updating pdoc
                 for (const [f, loc] of await getFiles('testdata')) {
                     if (f.isDirectory()) {
                         const sub = await fs.readdir(loc);
@@ -521,11 +549,8 @@ export class ProblemModel {
                     if (count > 5) continue;
                     await RecordModel.add(domainId, docId, operator, f.name.split('.')[1], await fs.readFile(loc, 'utf-8'), true);
                 }
-                if (process.env.HYDRO_CLI) {
-                    logger.info(`Imported problem ${pdoc.pid} (${pdoc.title})`);
-                } else {
-                    progress?.(`Imported problem ${pdoc.pid} (${pdoc.title})`);
-                }
+                const message = `${overridePid ? 'Updated' : 'Imported'} problem ${pdoc.pid} (${pdoc.title})`;
+                (process.env.HYDRO_CLI ? logger.info : progress)?.(message);
             }
         } finally {
             if (del) await fs.remove(tmpdir);