refactor!: rewrite streaming logic (#136)

googleapis · Jul 12, 2019 · 641d82d · 641d82d
1 parent 65ff625
commit 641d82d
Show file tree

Hide file tree

Showing 5 changed files with 446 additions and 507 deletions.
diff --git a/package.json b/package.json
@@ -52,9 +52,7 @@
   },
   "dependencies": {
     "arrify": "^2.0.0",
-    "extend": "^3.0.1",
-    "split-array-stream": "^2.0.0",
-    "stream-events": "^1.0.4"
+    "extend": "^3.0.1"
   },
   "engines": {
     "node": ">=8.10.0"

diff --git a/src/index.ts b/src/index.ts
@@ -20,83 +20,8 @@
 
 import arrify = require('arrify');
 import * as extend from 'extend';
-import {split} from 'split-array-stream';
-import {Transform, TransformOptions} from 'stream';
-import * as streamEvents from 'stream-events';
-
-export interface CreateLimiterOptions {
-  /**
-   * The maximum number of API calls to make.
-   */
-  maxApiCalls?: number;
-
-  /**
-   * Options to pass to the Stream constructor.
-   */
-  streamOptions?: TransformOptions;
-}
-
-export interface Limiter {
-  // tslint:disable-next-line no-any
-  makeRequest(...args: any[]): Transform | undefined;
-  stream: Transform;
-}
-
-export type ResourceStream<T> = {
-  addListener(event: 'data', listener: (data: T) => void): ResourceStream<T>;
-  emit(event: 'data', data: T): boolean;
-  on(event: 'data', listener: (data: T) => void): ResourceStream<T>;
-  once(event: 'data', listener: (data: T) => void): ResourceStream<T>;
-  prependListener(
-    event: 'data',
-    listener: (data: T) => void
-  ): ResourceStream<T>;
-  prependOnceListener(
-    event: 'data',
-    listener: (data: T) => void
-  ): ResourceStream<T>;
-  removeListener(event: 'data', listener: (data: T) => void): ResourceStream<T>;
-} & Transform;
-
-/**
- * Limit requests according to a `maxApiCalls` limit.
- *
- * @param {function} makeRequestFn - The function that will be called.
- * @param {object=} options - Configuration object.
- * @param {number} options.maxApiCalls - The maximum number of API calls to make.
- * @param {object} options.streamOptions - Options to pass to the Stream constructor.
- */
-export function createLimiter(
-  makeRequestFn: Function,
-  options?: CreateLimiterOptions
-): Limiter {
-  options = options || {};
-
-  const streamOptions = options.streamOptions || {};
-  streamOptions.objectMode = true;
-  const stream = streamEvents(new Transform(streamOptions)) as Transform;
-
-  let requestsMade = 0;
-  let requestsToMake = -1;
-
-  if (typeof options.maxApiCalls === 'number') {
-    requestsToMake = options.maxApiCalls!;
-  }
-
-  return {
-    // tslint:disable-next-line:no-any
-    makeRequest(...args: any[]) {
-      requestsMade++;
-      if (requestsToMake >= 0 && requestsMade > requestsToMake) {
-        stream.push(null);
-        return;
-      }
-      makeRequestFn.apply(null, args);
-      return stream;
-    },
-    stream,
-  };
-}
+import {TransformOptions} from 'stream';
+import {ResourceStream} from './resource-stream';
 
 export interface ParsedArguments extends TransformOptions {
   /**
@@ -196,7 +121,10 @@ export class Paginator {
     ): ResourceStream<T> {
       const parsedArguments = paginator.parseArguments_(args);
       const originalMethod = this[methodName + '_'] || this[methodName];
-      return paginator.runAsStream_(parsedArguments, originalMethod.bind(this));
+      return paginator.runAsStream_<T>(
+        parsedArguments,
+        originalMethod.bind(this)
+      );
     };
   }
 
@@ -327,52 +255,16 @@ export class Paginator {
    *     and returns `nextQuery` to receive more results.
    * @return {stream} - Readable object stream.
    */
-  runAsStream_(parsedArguments: ParsedArguments, originalMethod: Function) {
-    const query = parsedArguments.query;
-    let resultsToSend = parsedArguments.maxResults!;
-
-    const limiter = exports.createLimiter(makeRequest, {
-      maxApiCalls: parsedArguments.maxApiCalls,
-      streamOptions: parsedArguments.streamOptions,
-    });
-
-    const stream = limiter.stream;
-
-    stream.once('reading', () => {
-      limiter.makeRequest(query);
-    });
-
-    function makeRequest(query?: ParsedArguments | string) {
-      originalMethod(query, onResultSet);
-    }
-
-    // tslint:disable-next-line:no-any
-    function onResultSet(err: Error | null, results?: any[], nextQuery?: any) {
-      if (err) {
-        stream.destroy(err);
-        return;
-      }
-
-      if (resultsToSend >= 0 && results!.length > resultsToSend) {
-        results = results!.splice(0, resultsToSend);
-      }
-
-      resultsToSend -= results!.length;
-
-      split(results!, stream).then(streamEnded => {
-        if (streamEnded) {
-          return;
-        }
-        if (nextQuery && resultsToSend !== 0) {
-          limiter.makeRequest(nextQuery);
-          return;
-        }
-        stream.push(null);
-      });
-    }
-    return limiter.stream;
+  // tslint:disable-next-line:no-any
+  runAsStream_<T = any>(
+    parsedArguments: ParsedArguments,
+    originalMethod: Function
+  ): ResourceStream<T> {
+    return new ResourceStream<T>(parsedArguments, originalMethod);
   }
 }
 
 const paginator = new Paginator();
 export {paginator};
+
+export {ResourceStream};
diff --git a/src/resource-stream.ts b/src/resource-stream.ts
@@ -0,0 +1,101 @@
+/*!
+ * Copyright 2019 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {Transform} from 'stream';
+import {ParsedArguments} from './';
+
+interface ResourceEvents<T> {
+  addListener(event: 'data', listener: (data: T) => void): this;
+  emit(event: 'data', data: T): boolean;
+  on(event: 'data', listener: (data: T) => void): this;
+  once(event: 'data', listener: (data: T) => void): this;
+  prependListener(event: 'data', listener: (data: T) => void): this;
+  prependOnceListener(event: 'data', listener: (data: T) => void): this;
+  removeListener(event: 'data', listener: (data: T) => void): this;
+}
+
+export class ResourceStream<T> extends Transform implements ResourceEvents<T> {
+  _ended: boolean;
+  _maxApiCalls: number;
+  _nextQuery: {} | null;
+  _reading: boolean;
+  _requestFn: Function;
+  _requestsMade: number;
+  _resultsToSend: number;
+  constructor(args: ParsedArguments, requestFn: Function) {
+    const options = Object.assign({objectMode: true}, args.streamOptions);
+    super(options);
+
+    this._ended = false;
+    this._maxApiCalls = args.maxApiCalls === -1 ? Infinity : args.maxApiCalls!;
+    this._nextQuery = args.query!;
+    this._reading = false;
+    this._requestFn = requestFn;
+    this._requestsMade = 0;
+    this._resultsToSend = args.maxResults === -1 ? Infinity : args.maxResults!;
+  }
+  // tslint:disable-next-line:no-any
+  end(...args: any[]) {
+    this._ended = true;
+    return super.end(...args);
+  }
+  _read() {
+    if (this._reading) {
+      return;
+    }
+
+    this._reading = true;
+
+    this._requestFn(
+      this._nextQuery,
+      (err: Error | null, results: T[], nextQuery: {} | null) => {
+        if (err) {
+          this.destroy(err);
+          return;
+        }
+
+        this._nextQuery = nextQuery;
+
+        if (this._resultsToSend !== Infinity) {
+          results = results.splice(0, this._resultsToSend);
+          this._resultsToSend -= results.length;
+        }
+
+        let more = true;
+
+        for (const result of results) {
+          if (this._ended) {
+            break;
+          }
+          more = this.push(result);
+        }
+
+        const isFinished = !this._nextQuery || this._resultsToSend < 1;
+        const madeMaxCalls = ++this._requestsMade >= this._maxApiCalls;
+
+        if (isFinished || madeMaxCalls) {
+          this.end();
+        }
+
+        if (more && !this._ended) {
+          setImmediate(() => this._read());
+        }
+
+        this._reading = false;
+      }
+    );
+  }
+}