feat: add ctx for the static tpl of response (#4047)

* feat: add ctx for static tpl * test: add more test and remove not found
midwayjs · Sep 2, 2024 · 2abd7f5 · 2abd7f5
1 parent a6fa886
commit 2abd7f5
Show file tree

Hide file tree

Showing 7 changed files with 151 additions and 30 deletions.
diff --git a/packages/core/src/interface.ts b/packages/core/src/interface.ts
@@ -1238,11 +1238,11 @@ export interface ServerSendEventMessage {
   retry?: number;
 }
 
-export interface ServerStreamOptions {
-  tpl?: (data: unknown) => unknown;
+export interface ServerStreamOptions<CTX extends IMidwayContext> {
+  tpl?: (data: unknown, ctx: CTX) => unknown;
 }
 
-export interface ServerSendEventStreamOptions {
+export interface ServerSendEventStreamOptions<CTX extends IMidwayContext> {
   closeEvent?: string;
-  tpl?: (data: ServerSendEventMessage) => ServerSendEventMessage;
+  tpl?: (data: ServerSendEventMessage, ctx: CTX) => ServerSendEventMessage;
 }
diff --git a/packages/core/src/response/base.ts b/packages/core/src/response/base.ts
@@ -8,11 +8,19 @@ export class ServerResponse<CTX extends IMidwayContext = IMidwayContext> {
     this.ctx = ctx;
   }
 
-  static TEXT_TPL = (data: string, isSuccess: boolean): unknown => {
+  static TEXT_TPL = <CTX extends IMidwayContext>(
+    data: string,
+    isSuccess: boolean,
+    ctx: CTX
+  ): unknown => {
     return data;
   };
 
-  static JSON_TPL = (data: Record<any, any>, isSuccess: boolean): unknown => {
+  static JSON_TPL = <CTX extends IMidwayContext>(
+    data: Record<any, any>,
+    isSuccess: boolean,
+    ctx: CTX
+  ): unknown => {
     if (isSuccess) {
       return {
         success: 'true',
@@ -26,28 +34,35 @@ export class ServerResponse<CTX extends IMidwayContext = IMidwayContext> {
     }
   };
 
-  static BLOB_TPL = (data: Buffer, isSuccess: boolean): unknown => {
+  static BLOB_TPL = <CTX extends IMidwayContext>(
+    data: Buffer,
+    isSuccess: boolean,
+    ctx: CTX
+  ): unknown => {
     return data;
   };
 
   json(data: Record<any, any>) {
     return Object.getPrototypeOf(this).constructor.JSON_TPL(
       data,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 
   text(data: string) {
     return Object.getPrototypeOf(this).constructor.TEXT_TPL(
       data,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 
   blob(data: Buffer) {
     return Object.getPrototypeOf(this).constructor.BLOB_TPL(
       data,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 

diff --git a/packages/core/src/response/http.ts b/packages/core/src/response/http.ts
@@ -18,15 +18,30 @@ export class HttpServerResponse<
     super(ctx);
   }
 
-  static FILE_TPL = (data: Readable, isSuccess: boolean) => {
+  static FILE_TPL = <CTX extends IMidwayContext>(
+    data: Readable,
+    isSuccess: boolean,
+    ctx: CTX
+  ) => {
     return data;
   };
 
-  static SSE_TPL = (data: ServerSendEventMessage) => {
+  static SSE_TPL = <CTX extends IMidwayContext>(
+    data: ServerSendEventMessage,
+    ctx: CTX
+  ) => {
     return data;
   };
 
-  static STREAM_TPL = (data: unknown) => {
+  static STREAM_TPL = <CTX extends IMidwayContext>(data: unknown, ctx: CTX) => {
+    return data;
+  };
+
+  static HTML_TPL = <CTX extends IMidwayContext>(
+    data: string,
+    isSuccess: boolean,
+    ctx: CTX
+  ): unknown => {
     return data;
   };
 
@@ -55,15 +70,17 @@ export class HttpServerResponse<
     this.header('Content-Type', 'application/json');
     return Object.getPrototypeOf(this).constructor.JSON_TPL(
       data,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 
   text(data: string) {
     this.header('Content-Type', 'text/plain');
     return Object.getPrototypeOf(this).constructor.TEXT_TPL(
       data,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 
@@ -75,26 +92,48 @@ export class HttpServerResponse<
     );
     return Object.getPrototypeOf(this).constructor.FILE_TPL(
       typeof filePath === 'string' ? createReadStream(filePath) : filePath,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 
   blob(data: Buffer, mimeType?: string) {
     this.header('Content-Type', mimeType || 'application/octet-stream');
     return Object.getPrototypeOf(this).constructor.BLOB_TPL(
       data,
-      this.isSuccess
+      this.isSuccess,
+      this.ctx
     );
   }
 
-  sse(options: ServerSendEventStreamOptions = {}) {
+  html(data: string) {
+    this.header('Content-Type', 'text/html');
+    return Object.getPrototypeOf(this).constructor.HTML_TPL(
+      data,
+      this.isSuccess,
+      this.ctx
+    );
+  }
+
+  redirect(url: string, status = 302) {
+    this.status(status);
+    if (this.ctx.redirect) {
+      return this.ctx.redirect(url);
+    } else if (this.ctx.res.redirect) {
+      return this.ctx.res.redirect(url);
+    } else {
+      this.header('Location', url);
+    }
+  }
+
+  sse(options: ServerSendEventStreamOptions<CTX> = {}) {
     return new ServerSendEventStream(this.ctx, {
       tpl: Object.getPrototypeOf(this).constructor.SSE_TPL,
       ...options,
     });
   }
 
-  stream(options: ServerStreamOptions = {}) {
+  stream(options: ServerStreamOptions<CTX> = {}) {
     return new HttpStreamResponse(this.ctx, {
       tpl: Object.getPrototypeOf(this).constructor.STREAM_TPL,
       ...options,

diff --git a/packages/core/src/response/sse.ts b/packages/core/src/response/sse.ts
@@ -1,5 +1,5 @@
 import { Transform } from 'stream';
-import { ServerSendEventStreamOptions } from '../interface';
+import { IMidwayContext, ServerSendEventStreamOptions } from '../interface';
 
 interface MessageEvent {
   data?: string | object;
@@ -8,13 +8,15 @@ interface MessageEvent {
   retry?: number;
 }
 
-export class ServerSendEventStream extends Transform {
+export class ServerSendEventStream<
+  CTX extends IMidwayContext
+> extends Transform {
   private readonly ctx: any;
   private isActive = false;
   private readonly closeEvent: string;
-  private options: ServerSendEventStreamOptions;
+  private options: ServerSendEventStreamOptions<CTX>;
 
-  constructor(ctx, options: ServerSendEventStreamOptions = {}) {
+  constructor(ctx, options: ServerSendEventStreamOptions<CTX> = {}) {
     super({
       objectMode: true,
       ...options,
@@ -111,7 +113,7 @@ export class ServerSendEventStream extends Transform {
   }
 
   send(message: MessageEvent): void {
-    super.write(this.options.tpl(message));
+    super.write(this.options.tpl(message, this.ctx));
   }
 
   private handleClose() {

diff --git a/packages/core/src/response/stream.ts b/packages/core/src/response/stream.ts
@@ -1,12 +1,12 @@
 import { Transform } from 'stream';
-import { ServerStreamOptions } from '../interface';
+import { IMidwayContext, ServerStreamOptions } from '../interface';
 
-export class HttpStreamResponse extends Transform {
-  private ctx: any;
+export class HttpStreamResponse<CTX extends IMidwayContext> extends Transform {
+  private ctx: CTX & { res: any; req: any };
   private isActive = false;
-  private options: ServerStreamOptions;
+  private options: ServerStreamOptions<CTX>;
 
-  constructor(ctx, options: ServerStreamOptions = {}) {
+  constructor(ctx, options: ServerStreamOptions<CTX> = {}) {
     super({
       objectMode: true,
       ...options,
@@ -44,7 +44,7 @@ export class HttpStreamResponse extends Transform {
     if (!this.writable) {
       return;
     }
-    this.write(this.options.tpl(data));
+    this.write(this.options.tpl(data, this.ctx));
   }
 
   sendError(error) {

diff --git a/packages/core/test/response/base.test.ts b/packages/core/test/response/base.test.ts
@@ -73,4 +73,44 @@ describe('response/base.test.ts', () => {
     res = response.fail().blob(Buffer.from('abc'));
     expect(res).toEqual(Buffer.from('abc'));
   });
+
+  it("should test with ctx", () => {
+    class ServerResponseNew extends ServerResponse {}
+    const response = new ServerResponseNew({
+      data: {
+        text: '123',
+      }
+    } as any);
+    let res = response.success().text('abc');
+    expect(res).toEqual('abc');
+
+    ServerResponseNew.TEXT_TPL = (data: string, isSuccess, ctx: any) => {
+      return isSuccess ? {
+        bbb: 0,
+        data,
+        ...ctx.data,
+      } : {
+        bbb: -1,
+        message: data || 'fail',
+      };
+    }
+
+    ServerResponseNew.JSON_TPL = (data: Record<any, any>, isSuccess, ctx: any) => {
+      return isSuccess ? {
+        bbb: 2,
+        data,
+        ...ctx.data,
+      } : {
+        bbb: 2,
+        message: data || 'fail2',
+      };
+    }
+
+    res = response.success().text('abc');
+    expect(res).toEqual({ bbb: 0, data: 'abc', text: '123' });
+
+    // json
+    res = response.success().json({ a: 1 });
+    expect(res).toEqual({ bbb: 2, data: { a: 1 }, text: '123' });
+  });
 })
diff --git a/packages/core/test/response/http.test.ts b/packages/core/test/response/http.test.ts
@@ -519,5 +519,30 @@ describe('response/http.test.ts', () => {
       expect(content).toMatch(/@midwayjs\/core/);
       unlinkSync(join(__dirname, 'package.json'));
     });
+
+    it("should test html response", () => {
+      HttpServerResponse.HTML_TPL = (data, isSuccess, ctx) => {
+        return `<div>${data}</div>`;
+      }
+      const ctx = {
+        logger: console,
+        res: new ServerResponse({} as any),
+      } as any;
+      const res = new HttpServerResponse(ctx);
+      const html = res.html('hello');
+      expect(html).toEqual('<div>hello</div>');
+      expect(ctx.res.getHeader('Content-Type')).toBe('text/html');
+    });
+
+    it("should test redirect response", () => {
+      const ctx = {
+        logger: console,
+        res: new ServerResponse({} as any),
+      } as any;
+      const res = new HttpServerResponse(ctx);
+      res.redirect('https://www.baidu.com');
+      expect(ctx.res.getHeader('Location')).toBe('https://www.baidu.com');
+      expect(ctx.res.statusCode).toBe(302);
+    });
   });
 });