fix: should pass extra callback arguments back to consumer (#361)

googleapis · May 22, 2024 · cc5c48b · cc5c48b
1 parent a4b49e0
commit cc5c48b
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 6 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -221,18 +221,22 @@ export class Paginator {
       return originalMethod(query, callback);
     }
     const results = new Array<{}>();
+    let otherArgs: unknown[] = [];
     const promise = new Promise((resolve, reject) => {
-      paginator
-        .runAsStream_(parsedArguments, originalMethod)
+      const stream = paginator.runAsStream_(parsedArguments, originalMethod);
+      stream
         .on('error', reject)
         .on('data', (data: {}) => results.push(data))
-        .on('end', () => resolve(results));
+        .on('end', () => {
+          otherArgs = stream._otherArgs || [];
+          resolve(results);
+        });
     });
     if (!callback) {
-      return promise.then(results => [results]);
+      return promise.then(results => [results, ...otherArgs]);
     }
     promise.then(
-      results => callback(null, results),
+      results => callback(null, results, ...otherArgs),
       (err: Error) => callback(err)
     );
   }

diff --git a/src/resource-stream.ts b/src/resource-stream.ts
@@ -31,6 +31,7 @@ export class ResourceStream<T> extends Transform implements ResourceEvents<T> {
   _ended: boolean;
   _maxApiCalls: number;
   _nextQuery: {} | null;
+  _otherArgs: unknown[];
   _reading: boolean;
   _requestFn: Function;
   _requestsMade: number;
@@ -46,6 +47,7 @@ export class ResourceStream<T> extends Transform implements ResourceEvents<T> {
     this._requestFn = requestFn;
     this._requestsMade = 0;
     this._resultsToSend = args.maxResults === -1 ? Infinity : args.maxResults!;
+    this._otherArgs = [];
   }
   /* eslint-disable  @typescript-eslint/no-explicit-any */
   end(
@@ -67,12 +69,18 @@ export class ResourceStream<T> extends Transform implements ResourceEvents<T> {
     try {
       this._requestFn(
         this._nextQuery,
-        (err: Error | null, results: T[], nextQuery: {} | null) => {
+        (
+          err: Error | null,
+          results: T[],
+          nextQuery: {} | null,
+          ...otherArgs: []
+        ) => {
           if (err) {
             this.destroy(err);
             return;
           }
 
+          this._otherArgs = otherArgs;
           this._nextQuery = nextQuery;
 
           if (this._resultsToSend !== Infinity) {

diff --git a/test/index.ts b/test/index.ts
@@ -382,6 +382,38 @@ describe('paginator', () => {
 
           paginator.run_(parsedArguments, util.noop);
         });
+
+        it('should return all results and extra args', done => {
+          const results = [{a: 1}, {b: 2}, {c: 3}];
+          const args: any[] = [{msg: 'OK'}, 10];
+
+          const parsedArguments = {
+            autoPaginate: true,
+            callback(
+              err: Error,
+              results_: {},
+              fakeRes: {},
+              anotherArg: number
+            ) {
+              assert.deepStrictEqual(results_, results);
+              assert.deepStrictEqual(fakeRes, {msg: 'OK'});
+              assert.deepStrictEqual(anotherArg, 10);
+              done();
+            },
+          };
+
+          sandbox.stub(paginator, 'runAsStream_').callsFake(() => {
+            const stream = createFakeStream();
+            setImmediate(() => {
+              results.forEach(result => stream.push(result));
+              stream.push(null);
+              stream._otherArgs = args;
+            });
+            return stream;
+          });
+
+          paginator.run_(parsedArguments, util.noop);
+        });
       });
 
       describe('original method is promise based', () => {
@@ -434,6 +466,29 @@ describe('paginator', () => {
               assert.deepStrictEqual(results_, results)
             );
         });
+
+        it('should resolve with all results and extra args', () => {
+          const results = [{a: 1}, {b: 2}, {c: 3}];
+          const args: any[] = [{msg: 'OK'}, 10];
+
+          sandbox.stub(paginator, 'runAsStream_').callsFake(() => {
+            const stream = createFakeStream();
+            setImmediate(() => {
+              results.forEach(result => stream.push(result));
+              stream.push(null);
+              stream._otherArgs = args;
+            });
+            return stream;
+          });
+
+          paginator
+            .run_(parsedArguments, util.noop)
+            .then(([results_, fakeRes, anotherArg]: unknown[]) => {
+              assert.deepStrictEqual(results_, results);
+              assert.deepEqual(fakeRes, {msg: 'OK'});
+              assert.deepEqual(anotherArg, 10);
+            });
+        });
       });
     });
 

diff --git a/test/resource-stream.ts b/test/resource-stream.ts
@@ -150,6 +150,17 @@ describe('ResourceStream', () => {
       assert.strictEqual(stream._nextQuery, fakeQuery);
     });
 
+    it('should cache the rest of the callback arguments', () => {
+      const fakeRes = {status: 'OK'};
+      const anotherArg = 10;
+
+      stream._read();
+      const callback = requestSpy.lastCall.args[1];
+      callback(null, [], {}, fakeRes, anotherArg);
+
+      assert.deepStrictEqual(stream._otherArgs, [fakeRes, anotherArg]);
+    });
+
     it('should adjust the results to send counter', () => {
       const maxResults = 100;
       const results = [{}, {}];