Browse Source

model-server: refactored result writer

David Sehnal 5 years ago
parent
commit
6eee3e8368

+ 6 - 38
src/servers/model/server/api-local.ts

@@ -4,16 +4,14 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import * as fs from 'fs';
-import * as path from 'path';
-import { JobManager, Job, JobEntry } from './jobs';
 import { ConsoleLogger } from '../../../mol-util/console-logger';
-import { resolveJob } from './query';
-import { StructureCache } from './structure-wrapper';
 import { now } from '../../../mol-util/now';
 import { PerformanceMonitor } from '../../../mol-util/performance-monitor';
+import { FileResultWriter } from '../utils/writer';
 import { QueryName } from './api';
-import { makeDir } from '../../../mol-util/make-dir';
+import { Job, JobEntry, JobManager } from './jobs';
+import { resolveJob } from './query';
+import { StructureCache } from './structure-wrapper';
 
 export type LocalInput = {
     input: string,
@@ -39,6 +37,7 @@ export async function runLocal(input: LocalInput) {
                 queryParams: job.params || { },
                 modelNums: job.modelNums,
             })],
+            writer: new FileResultWriter(job.output),
             options: {
                 outputFilename: job.output,
                 binary
@@ -55,7 +54,7 @@ export async function runLocal(input: LocalInput) {
     while (job) {
         try {
             const encoder = await resolveJob(job);
-            const writer = wrapFileToWriter(job.outputFilename!);
+            const writer = job.writer;
             encoder.writeTo(writer);
             writer.end();
             ConsoleLogger.logId(job.id, 'Query', 'Written.');
@@ -75,35 +74,4 @@ export async function runLocal(input: LocalInput) {
 
     ConsoleLogger.log('Progress', `Done in ${PerformanceMonitor.format(now() - started)}.`);
     StructureCache.expireAll();
-}
-
-export function wrapFileToWriter(fn: string) {
-    const w = {
-        open(this: any) {
-            if (this.opened) return;
-            makeDir(path.dirname(fn));
-            this.file = fs.openSync(fn, 'w');
-            this.opened = true;
-        },
-        writeBinary(this: any, data: Uint8Array) {
-            this.open();
-            fs.writeSync(this.file, Buffer.from(data.buffer));
-            return true;
-        },
-        writeString(this: any, data: string) {
-            this.open();
-            fs.writeSync(this.file, data);
-            return true;
-        },
-        end(this: any) {
-            if (!this.opened || this.ended) return;
-            fs.close(this.file, function () { });
-            this.ended = true;
-        },
-        file: 0,
-        ended: false,
-        opened: false
-    };
-
-    return w;
 }

+ 12 - 46
src/servers/model/server/api-web.ts

@@ -17,67 +17,24 @@ import { QueryDefinition, normalizeRestQueryParams, normalizeRestCommonParams, Q
 import { getApiSchema, shortcutIconLink } from './api-schema';
 import { swaggerUiAssetsHandler, swaggerUiIndexHandler } from '../../common/swagger-ui';
 import { MultipleQuerySpec } from './api-web-multiple';
+import { SimpleResponseResultWriter, WebResutlWriter } from '../utils/writer';
 
 function makePath(p: string) {
     return Config.apiPrefix + '/' + p;
 }
 
-function wrapResponse(fn: string, res: express.Response) {
-    const w = {
-        doError(this: any, code = 404, message = 'Not Found.') {
-            if (!this.headerWritten) {
-                res.status(code).send(message);
-                this.headerWritten = true;
-            }
-            this.end();
-        },
-        writeHeader(this: any, binary: boolean) {
-            if (this.headerWritten) return;
-            res.writeHead(200, {
-                'Content-Type': binary ? 'application/octet-stream' : 'text/plain; charset=utf-8',
-                'Access-Control-Allow-Origin': '*',
-                'Access-Control-Allow-Headers': 'X-Requested-With',
-                'Content-Disposition': `inline; filename="${fn}"`
-            });
-            this.headerWritten = true;
-        },
-        writeBinary(this: any, data: Uint8Array) {
-            if (!this.headerWritten) this.writeHeader(true);
-            return res.write(Buffer.from(data.buffer));
-        },
-        writeString(this: any, data: string) {
-            if (!this.headerWritten) this.writeHeader(false);
-            return res.write(data);
-        },
-        end(this: any) {
-            if (this.ended) return;
-            res.end();
-            this.ended = true;
-        },
-        ended: false,
-        headerWritten: false
-    };
-
-    return w;
-}
-
 const responseMap = new Map<UUID, express.Response>();
 
 async function processNextJob() {
     if (!JobManager.hasNext()) return;
 
     const job = JobManager.getNext();
-    const response = responseMap.get(job.id)!;
     responseMap.delete(job.id);
-
-    const filenameBase = job.entries.length === 1
-        ? `${job.entries[0].entryId}_${job.entries[0].queryDefinition.name.replace(/\s/g, '_')}`
-        : `result`;
-    const writer = wrapResponse(job.responseFormat.isBinary ? `${filenameBase}.bcif` : `${filenameBase}.cif`, response);
+    const writer = job.writer as WebResutlWriter;
 
     try {
         const encoder = await resolveJob(job);
-        writer.writeHeader(job.responseFormat.isBinary);
+        writer.writeHeader();
         encoder.writeTo(writer);
     } catch (e) {
         ConsoleLogger.errorId(job.id, '' + e);
@@ -89,6 +46,13 @@ async function processNextJob() {
     }
 }
 
+export function createResultWriter(response: express.Response, isBinary: boolean, entryId?: string, queryName?: string) {
+    const filenameBase = entryId && queryName
+        ? `${entryId}_${queryName.replace(/\s/g, '_')}`
+        : `result`;
+    return new SimpleResponseResultWriter(isBinary ? `${filenameBase}.bcif` : `${filenameBase}.cif`, response, isBinary);
+}
+
 function mapQuery(app: express.Express, queryName: string, queryDefinition: QueryDefinition) {
     function createJob(queryParams: any, req: express.Request, res: express.Response) {
         const entryId = req.params.id;
@@ -101,6 +65,7 @@ function mapQuery(app: express.Express, queryName: string, queryDefinition: Quer
                 queryParams,
                 modelNums: commonParams.model_nums
             })],
+            writer: createResultWriter(res, commonParams.encoding === 'bcif', entryId, queryName),
             options: { binary: commonParams.encoding === 'bcif' }
         });
         responseMap.set(jobId, res);
@@ -162,6 +127,7 @@ function createMultiJob(spec: MultipleQuerySpec, res: express.Response) {
             queryParams: q.params || { },
             modelNums: q.model_nums
         })),
+        writer: createResultWriter(res, spec.encoding?.toLowerCase() === 'bcif'),
         options: { binary: spec.encoding?.toLowerCase() === 'bcif' }
     });
     responseMap.set(jobId, res);

+ 6 - 1
src/servers/model/server/jobs.ts

@@ -7,6 +7,7 @@
 import { UUID } from '../../../mol-util';
 import { getQueryByName, QueryDefinition, QueryName, QueryParams } from './api';
 import { LinkedList } from '../../../mol-data/generic';
+import { ResultWriter } from '../utils/writer';
 
 export interface ResponseFormat {
     isBinary: boolean
@@ -19,7 +20,9 @@ export interface Job {
     entries: JobEntry[],
 
     responseFormat: ResponseFormat,
-    outputFilename?: string
+    outputFilename?: string,
+
+    writer: ResultWriter
 }
 
 export interface JobEntry {
@@ -61,6 +64,7 @@ export function JobEntry<Name extends QueryName>(definition: JobEntryDefinition<
 
 export interface JobDefinition {
     entries: JobEntry[],
+    writer: ResultWriter,
     options?: { outputFilename?: string, binary?: boolean }
 }
 
@@ -69,6 +73,7 @@ export function createJob(definition: JobDefinition): Job {
         id: UUID.create22(),
         datetime_utc: `${new Date().toISOString().replace(/T/, ' ').replace(/\..+/, '')}`,
         entries: definition.entries,
+        writer: definition.writer,
         responseFormat: { isBinary: !!(definition.options && definition.options.binary) },
         outputFilename: definition.options && definition.options.outputFilename
     };

+ 111 - 0
src/servers/model/utils/writer.ts

@@ -0,0 +1,111 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import * as express from 'express';
+import * as fs from 'fs';
+import * as path from 'path';
+import { makeDir } from '../../../mol-util/make-dir';
+
+export interface ResultWriter {
+    beginEntry(name: string): void,
+    endEntry(): void,
+
+    writeBinary(data: Uint8Array): boolean,
+    writeString(data: string): boolean,
+    end(): void
+}
+
+export interface WebResutlWriter extends ResultWriter {
+    doError(code?: number, message?: string): void,
+    writeHeader(): void
+}
+
+export class SimpleResponseResultWriter implements WebResutlWriter {
+    private ended = false;
+
+    beginEntry(name: string) {
+        throw new Error('Not supported');
+    }
+
+    endEntry() {
+        throw new Error('Not supported');
+    }
+
+    doError(code = 404, message = 'Not Found.') {
+        this.res.status(code).send(message);
+        this.end();
+    }
+
+    writeHeader() {
+        this.res.writeHead(200, {
+            'Content-Type': this.isBinary ? 'application/octet-stream' : 'text/plain; charset=utf-8',
+            'Access-Control-Allow-Origin': '*',
+            'Access-Control-Allow-Headers': 'X-Requested-With',
+            'Content-Disposition': `inline; filename="${this.fn}"`
+        });
+    }
+
+    writeBinary(data: Uint8Array) {
+        return this.res.write(Buffer.from(data.buffer));
+    }
+
+    writeString(this: any, data: string) {
+        return this.res.write(data);
+    }
+
+    end() {
+        if (this.ended) return;
+        this.res.end();
+        this.ended = true;
+    }
+
+    constructor(private fn: string, private res: express.Response, private isBinary: boolean) {
+
+    }
+}
+
+export class FileResultWriter implements ResultWriter {
+    private file = 0;
+    private ended = false;
+    private opened = false;
+
+    beginEntry(name: string) {
+        throw new Error('Not supported');
+    }
+
+    endEntry() {
+        throw new Error('Not supported');
+    }
+
+    open() {
+        if (this.opened) return;
+        makeDir(path.dirname(this.fn));
+        this.file = fs.openSync(this.fn, 'w');
+        this.opened = true;
+    }
+
+    writeBinary(data: Uint8Array) {
+        this.open();
+        fs.writeSync(this.file, Buffer.from(data.buffer));
+        return true;
+    }
+
+    writeString(data: string) {
+        this.open();
+        fs.writeSync(this.file, data);
+        return true;
+    }
+
+    end() {
+        if (!this.opened || this.ended) return;
+        fs.close(this.file, function () { });
+        this.ended = true;
+    }
+
+    constructor(private fn: string) {
+
+    }
+}