refactor: replaces job#getQueryResults with table#getRows in query method

src/index.ts

@@ -1016,7 +1016,6 @@ export class BigQuery extends common.Service {
     if ((!options || !options.query) && !options.pageToken) {
       throw new Error('A SQL query string is required.');
     }
-
     // tslint:disable-next-line no-any
     const query: any = extend(
       true,
@@ -1457,8 +1456,7 @@ export class BigQuery extends common.Service {
     return new Job(this, id, options);
   }
 
-  query(query: string, options?: QueryOptions): Promise<QueryRowsResponse>;
-  query(query: Query, options?: QueryOptions): Promise<SimpleQueryRowsResponse>;
+  query(query: string, options?: QueryOptions): Promise<RowsResponse>;
   query(
     query: string,
     options: QueryOptions,
@@ -1565,9 +1563,10 @@ export class BigQuery extends common.Service {
     optionsOrCallback?:
       | QueryOptions
       | SimpleQueryRowsCallback
+      | RowsCallback
       | QueryRowsCallback,
-    cb?: SimpleQueryRowsCallback | QueryRowsCallback
-  ): void | Promise<SimpleQueryRowsResponse> | Promise<QueryRowsResponse> {
+    cb?: SimpleQueryRowsCallback | RowsCallback
+  ): void | Promise<RowsResponse> {
     let options =
       typeof optionsOrCallback === 'object' ? optionsOrCallback : {};
     const callback =
@@ -1583,8 +1582,26 @@ export class BigQuery extends common.Service {
       }
       // The Job is important for the `queryAsStream_` method, so a new query
       // isn't created each time results are polled for.
-      options = extend({job}, options);
-      job!.getQueryResults(options, callback as QueryRowsCallback);
+      job!
+        .on('error', err => callback!(err))
+        .on('complete', () => {
+          options = extend({job}, options);
+
+          if (
+            typeof query === 'string' &&
+            (query.includes('CREATE MODEL') || query.includes('REPLACE MODEL'))
+          ) {
+            job!.getQueryResults(options, callback as QueryRowsCallback);
+          } else {
+            const datasetId = job!.metadata.configuration.query.destinationTable
+              .datasetId;
+            const tableId = job!.metadata.configuration.query.destinationTable
+              .tableId;
+            const dataset = this.dataset(datasetId);
+            const table = dataset.table(tableId);
+            table.getRows(options, callback as RowsCallback);
+          }
+        });
     });
   }
 
@@ -1596,7 +1613,7 @@ export class BigQuery extends common.Service {
    */
   queryAsStream_(query: Query, callback?: SimpleQueryRowsCallback) {
     if (query.job) {
-      query.job.getQueryResults(query, callback as QueryRowsCallback);
+      query.job.getQueryResults(query, callback as RowsCallback);
       return;
     }
     this.query(query, {autoPaginate: false}, callback);

test/index.ts

@@ -29,6 +29,7 @@ import * as extend from 'extend';
 import * as proxyquire from 'proxyquire';
 import * as sinon from 'sinon';
 import * as uuid from 'uuid';
+import {EventEmitter} from 'events';
 
 import {BigQueryDate, Dataset, Job, Query, Table} from '../src';
 import {JobOptions} from '../src/job';
@@ -66,52 +67,6 @@ const fakeUtil = extend({}, util, {
 });
 const originalFakeUtil = extend(true, {}, fakeUtil);
 
-class FakeDataset {
-  calledWith_: IArguments;
-  constructor() {
-    this.calledWith_ = arguments;
-  }
-}
-
-class FakeTable extends Table {
-  constructor(a: Dataset, b: string) {
-    super(a, b);
-  }
-}
-
-class FakeJob {
-  calledWith_: IArguments;
-  constructor() {
-    this.calledWith_ = arguments;
-  }
-}
-
-let extended = false;
-const fakePaginator = {
-  paginator: {
-    extend: (c: Function, methods: string[]) => {
-      if (c.name !== 'BigQuery') {
-        return;
-      }
-      methods = arrify(methods);
-      assert.strictEqual(c.name, 'BigQuery');
-      assert.deepStrictEqual(methods, ['getDatasets', 'getJobs']);
-      extended = true;
-    },
-    streamify: (methodName: string) => {
-      return methodName;
-    },
-  },
-};
-
-class FakeService extends Service {
-  calledWith_: IArguments;
-  constructor(config: ServiceConfig, options: ServiceOptions) {
-    super(config, options);
-    this.calledWith_ = arguments;
-  }
-}
-
 const sandbox = sinon.createSandbox();
 afterEach(() => sandbox.restore());
 
@@ -127,6 +82,57 @@ describe('BigQuery', () => {
   // tslint:disable-next-line no-any
   let bq: any;
 
+  class FakeTable extends Table {
+    constructor(a: Dataset, b: string) {
+      super(a, b);
+    }
+  }
+
+  class FakeDataset extends Dataset {
+    calledWith_: IArguments;
+    constructor() {
+      super(bq, '1');
+      this.calledWith_ = arguments;
+    }
+
+    table(id: string): FakeTable {
+      return new FakeTable(this, id);
+    }
+  }
+
+  class FakeJob {
+    calledWith_: IArguments;
+    constructor() {
+      this.calledWith_ = arguments;
+    }
+  }
+
+  let extended = false;
+  const fakePaginator = {
+    paginator: {
+      extend: (c: Function, methods: string[]) => {
+        if (c.name !== 'BigQuery') {
+          return;
+        }
+        methods = arrify(methods);
+        assert.strictEqual(c.name, 'BigQuery');
+        assert.deepStrictEqual(methods, ['getDatasets', 'getJobs']);
+        extended = true;
+      },
+      streamify: (methodName: string) => {
+        return methodName;
+      },
+    },
+  };
+
+  class FakeService extends Service {
+    calledWith_: IArguments;
+    constructor(config: ServiceConfig, options: ServiceOptions) {
+      super(config, options);
+      this.calledWith_ = arguments;
+    }
+  }
+
   before(() => {
     BigQuery = proxyquire('../src', {
       uuid: fakeUuid,
@@ -1747,8 +1753,7 @@ describe('BigQuery', () => {
     const FAKE_ROWS = [{}, {}, {}];
     const FAKE_RESPONSE = {};
     const QUERY_STRING = 'SELECT * FROM [dataset.table]';
-
-    it('should return any errors from createQueryJob', done => {
+    const FAKE_JOB = it('should return any errors from createQueryJob', done => {
       const error = new Error('err');
 
       bq.createQueryJob = (query: {}, callback: Function) => {
@@ -1782,55 +1787,48 @@ describe('BigQuery', () => {
       });
     });
 
-    it('should call job#getQueryResults', done => {
-      const fakeJob = {
-        getQueryResults: (options: {}, callback: Function) => {
-          callback(null, FAKE_ROWS, FAKE_RESPONSE);
+    it('should call table#getRows', done => {
+      const metadata = {
+        configuration: {
+          query: {
+            destinationTable: {
+              datasetId: '1',
+              tableId: '1',
+            },
+          },
         },
       };
 
+      const fakeJob = new EventEmitter();
+      fakeJob.emit('complete', metadata);
+
       bq.createQueryJob = (query: {}, callback: Function) => {
         callback(null, fakeJob, FAKE_RESPONSE);
+        done();
       };
 
-      bq.query(QUERY_STRING, (err: Error, rows: {}, resp: {}) => {
+      const finalCallback = (err: Error | null, rows: {}, resp: {}) => {
         assert.ifError(err);
         assert.strictEqual(rows, FAKE_ROWS);
         assert.strictEqual(resp, FAKE_RESPONSE);
-        done();
-      });
-    });
-
-    it('should assign Job on the options', done => {
-      const fakeJob = {
-        getQueryResults: (options: {}, callback: Function) => {
-          assert.deepStrictEqual(options, {job: fakeJob});
-          done();
-        },
-      };
-
-      bq.createQueryJob = (query: {}, callback: Function) => {
-        callback(null, fakeJob, FAKE_RESPONSE);
       };
 
-      bq.query(QUERY_STRING, assert.ifError);
-    });
-
-    it('should optionally accept options', done => {
-      const fakeOptions = {};
-      const fakeJob = {
-        getQueryResults: (options: {}) => {
-          assert.notStrictEqual(options, fakeOptions);
+      const fakeTable = {
+        getRows(options: {}, cb: Function) {
           assert.deepStrictEqual(options, {job: fakeJob});
-          done();
+          assert.strictEqual(cb, finalCallback);
+          finalCallback(null, FAKE_ROWS, FAKE_RESPONSE);
         },
       };
 
-      bq.createQueryJob = (query: {}, callback: Function) => {
-        callback(null, fakeJob, FAKE_RESPONSE);
+      const fakeDataset = {
+        table(id: string) {
+          return fakeTable;
+        },
       };
 
-      bq.query(QUERY_STRING, fakeOptions, assert.ifError);
+      bq.dataset = (id: string) => fakeDataset;
+      bq.query(QUERY_STRING, finalCallback);
     });
   });