Change PipelineTransform to PipelineTransformationEntrypoint class implementation

Author: hhoughgg

src/cloudflare/internal/pipeline-transform.ts

@@ -1,3 +1,4 @@
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
 // Copyright (c) 2024 Cloudflare, Inc.
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
@@ -57,26 +58,39 @@ enum Format {
   JSON_STREAM = 'json_stream', // jsonl
 }
 
-export class PipelineTransformImpl extends entrypoints.WorkerEntrypoint {
+type PipelineBatchMetadata = {
+  pipelineId: string;
+  pipelineName: string;
+};
+
+type PipelineRecord = Record<string, unknown>;
+
+export class PipelineTransformImpl<
+  I extends PipelineRecord,
+  O extends PipelineRecord,
+> extends entrypoints.WorkerEntrypoint {
   #batch?: Batch;
   #initalized: boolean = false;
 
-  // stub overridden on the sub class
+  // stub overridden on the subclass
   // eslint-disable-next-line @typescript-eslint/require-await
-  public async transformJson(_data: object[]): Promise<object[]> {
+  public async run(
+    _records: I[],
+    _metadata: PipelineBatchMetadata
+  ): Promise<O[]> {
     throw new Error('should be implemented by parent');
   }
 
-  // called by the dispatcher which then calls the subclass methods
+  // called by the dispatcher to validate that run is properly implemented by the subclass
   // @ts-expect-error thinks ping is never used
   private _ping(): Promise<void> {
     // making sure the function was overridden by an implementing subclass
-    if (this.transformJson !== PipelineTransformImpl.prototype.transformJson) {
+    if (this.run !== PipelineTransformImpl.prototype.run) {
       return Promise.resolve();
     } else {
       return Promise.reject(
         new Error(
-          'the transformJson method must be overridden by the PipelineTransform subclass'
+          'the run method must be overridden by the PipelineTransformationEntrypoint subclass'
         )
       );
     }
@@ -85,55 +99,59 @@ export class PipelineTransformImpl extends entrypoints.WorkerEntrypoint {
   // called by the dispatcher which then calls the subclass methods
   // the reason this is typescript private and not javascript private is that this must be
   // able to be called by the dispatcher but should not be called by the class implementer
-  // @ts-expect-error _transform is called by rpc
-  private async _transform(batch: Batch): Promise<JsonStream> {
+  // @ts-expect-error _run is called by rpc
+  private async _run(
+    batch: Batch,
+    metadata: PipelineBatchMetadata
+  ): Promise<JsonStream> {
     if (this.#initalized) {
       throw new Error('pipeline entrypoint has already been initialized');
     }
 
     this.#batch = batch;
     this.#initalized = true;
 
-    switch (this.#batch.format) {
-      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-      case Format.JSON_STREAM: {
-        const data = await this.#readJsonStream();
-        const transformed = await this.transformJson(data);
-        return this.#sendJson(transformed);
-      }
-      default:
-        throw new Error('unsupported batch format');
+    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+    if (this.#batch.format === Format.JSON_STREAM) {
+      const records: I[] = await this.#readJsonStream();
+      const transformed = await this.run(records, metadata);
+      return this.#sendJson(transformed);
+    } else {
+      throw new Error(
+        'PipelineTransformationEntrypoint run supports only the JSON_STREAM batch format'
+      );
     }
   }
 
-  async #readJsonStream(): Promise<object[]> {
+  async #readJsonStream(): Promise<I[]> {
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
     if (this.#batch!.format !== Format.JSON_STREAM) {
+      // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
       throw new Error(`expected JSON_STREAM not ${this.#batch!.format}`);
     }
 
     const batch = this.#batch!.data as ReadableStream<Uint8Array>;
     const decoder = batch.pipeThrough(new TextDecoderStream());
 
-    const data: object[] = [];
+    const data: I[] = [];
     for await (const line of readLines(decoder)) {
-      data.push(JSON.parse(line) as object);
+      data.push(JSON.parse(line) as I);
     }
 
     return data;
   }
 
-  #sendJson(data: object[]): JsonStream {
-    if (!(data instanceof Array)) {
-      throw new Error('transformJson must return an array of objects');
+  #sendJson(records: O[]): JsonStream {
+    if (!(records instanceof Array)) {
+      throw new Error('transformations must return an array of PipelineRecord');
     }
 
     let written = 0;
     const encoder = new TextEncoder();
     const readable = new ReadableStream<Uint8Array>({
       start(controller): void {
-        for (const obj of data) {
-          const encoded = encoder.encode(`${JSON.stringify(obj)}\n`);
+        for (const record of records) {
+          const encoded = encoder.encode(`${JSON.stringify(record)}\n`);
           written += encoded.length;
           controller.enqueue(encoded);
         }
@@ -149,7 +167,7 @@ export class PipelineTransformImpl extends entrypoints.WorkerEntrypoint {
       format: Format.JSON_STREAM,
       size: {
         bytes: written,
-        rows: data.length,
+        rows: records.length,
       },
       data: readable,
     };

src/cloudflare/internal/test/pipeline-transform/transform-test.js

@@ -4,22 +4,22 @@
 //     https://opensource.org/licenses/Apache-2.0
 
 import assert from 'node:assert';
-import { PipelineTransform } from 'cloudflare:pipeline-transform';
+import { PipelineTransformationEntrypoint } from 'cloudflare:pipeline-transform';
 
 // this is how "Pipeline" would be implemented by the user
-const customTransform = class MyEntrypoint extends PipelineTransform {
+const customTransform = class MyEntrypoint extends PipelineTransformationEntrypoint {
   /**
    * @param {any} batch
    * @override
    */
-  async transformJson(batch) {
-    for (const obj of batch) {
-      obj.dispatcher = 'was here!';
+  async run(records, _) {
+    for (const record of records) {
+      record.dispatcher = 'was here!';
       await new Promise((resolve) => setTimeout(resolve, 50));
-      obj.wait = 'happened!';
+      record.wait = 'happened!';
     }
 
-    return batch;
+    return records;
   }
 };
 
@@ -52,11 +52,11 @@ export const tests = {
   async test(ctr, env, ctx) {
     {
       // should fail dispatcher test call when PipelineTransform class not extended
-      const transformer = new PipelineTransform(ctx, env);
+      const transformer = new PipelineTransformationEntrypoint(ctx, env);
       await assert.rejects(transformer._ping(), (err) => {
         assert.strictEqual(
           err.message,
-          'the transformJson method must be overridden by the PipelineTransform subclass'
+          'the run method must be overridden by the PipelineTransformationEntrypoint subclass'
         );
         return true;
       });
@@ -73,7 +73,10 @@ export const tests = {
       const transformer = new customTransform(ctx, env);
       const batch = newBatch();
 
-      const result = await transformer._transform(batch);
+      const result = await transformer._run(batch, {
+        id: 'abc',
+        name: 'mypipeline',
+      });
       assert.equal(true, result.data instanceof ReadableStream);
 
       const reader = result.data

src/cloudflare/pipeline-transform.ts

@@ -1,3 +1,3 @@
 import { PipelineTransformImpl } from 'cloudflare-internal:pipeline-transform';
 
-export const PipelineTransform = PipelineTransformImpl;
+export const PipelineTransformationEntrypoint = PipelineTransformImpl;

types/defines/pipeline-transform.d.ts

@@ -2,12 +2,20 @@
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
 
-export abstract class PipelineTransform {
+export type PipelineRecord = Record<string, unknown>
+
+export type PipelineBatchMetadata = {
+  pipelineId: string;
+  pipelineName: string;
+};
+
+export abstract class PipelineTransformationEntrypoint<I extends PipelineRecord, O extends PipelineRecord> {
   /**
-   * transformJson recieves an array of javascript objects which can be
+   * run recieves an array of PipelineRecord which can be
    * mutated and returned to the pipeline
-   * @param data The data to be mutated
-   * @returns A promise containing the mutated data
+   * @param records Incoming records from the pipeline to be transformed
+   * @param metadata Information about the specific pipeline calling the transformation entrypoint
+   * @returns A promise containing the transformed PipelineRecord array
    */
-  public transformJson(data: object[]): Promise<object[]>;
+  public run(records: I[], metadata: PipelineBatchMetadata): Promise<O[]>;
 }

types/defines/pipelines.d.ts

@@ -2,12 +2,13 @@
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
 
-export interface Pipeline {
+import { PipelineRecord } from "./pipeline-transform";
+
+export interface Pipeline<T extends PipelineRecord> {
   /**
-   * send takes an array of javascript objects which are
-   * then received by the pipeline for processing
+   * The Pipeline interface represents the type of a binding to a Pipeline
    *
-   * @param data The data to be sent
+   * @param records The records to send to the pipeline
    */
-  send(data: object[]): Promise<void>
+  send(records: T[]): Promise<void>
 }