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 =
@@ -1584,7 +1583,25 @@ 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);
+      if (
+        typeof query === 'string' &&
+        (query.includes('CREATE MODEL') || query.includes('REPLACE MODEL'))
+      ) {
+        job!.getQueryResults(options, callback as QueryRowsCallback);
+        return;
+      }
+
+      job!
+        .on('error', err => callback!(err))
+        .on('complete', () => {
+          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);
+        });
     });
   }
 

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,
@@ -1748,21 +1754,6 @@ describe('BigQuery', () => {
     const FAKE_RESPONSE = {};
     const QUERY_STRING = 'SELECT * FROM [dataset.table]';
 
-    it('should return any errors from createQueryJob', done => {
-      const error = new Error('err');
-
-      bq.createQueryJob = (query: {}, callback: Function) => {
-        callback(error, null, FAKE_RESPONSE);
-      };
-
-      bq.query(QUERY_STRING, (err: Error, rows: {}, resp: {}) => {
-        assert.strictEqual(err, error);
-        assert.strictEqual(rows, null);
-        assert.strictEqual(resp, FAKE_RESPONSE);
-        done();
-      });
-    });
-
     it('should exit early if dryRun is set', done => {
       const options = {
         query: QUERY_STRING,
@@ -1782,55 +1773,78 @@ 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 TABLE_ID = 'bq_table';
+      const DATASET_ID = 'bq_dataset';
+      const metadata = {
+        configuration: {
+          query: {
+            destinationTable: {
+              datasetId: DATASET_ID,
+              tableId: TABLE_ID,
+            },
+          },
         },
       };
 
+      const fakeJob = new EventEmitter();
+      // tslint:disable-next-line: no-any
+      (fakeJob as any).metadata = metadata;
+
       bq.createQueryJob = (query: {}, callback: Function) => {
         callback(null, fakeJob, FAKE_RESPONSE);
       };
 
-      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) => {
+      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) {
+          assert.strictEqual(id, TABLE_ID);
+          return fakeTable;
+        },
       };
 
-      bq.query(QUERY_STRING, assert.ifError);
+      bq.dataset = (id: string) => {
+        assert.strictEqual(id, DATASET_ID);
+        return fakeDataset;
+      };
+
+      bq.query(QUERY_STRING, finalCallback);
+      fakeJob.emit('complete', metadata);
     });
 
-    it('should optionally accept options', done => {
-      const fakeOptions = {};
+    it('should call job#getQueryResults for model query', done => {
+      const MODEL_QUERY_STRING = `CREATE OR REPLACE MODEL \`dataset.model\``;
+
       const fakeJob = {
-        getQueryResults: (options: {}) => {
-          assert.notStrictEqual(options, fakeOptions);
-          assert.deepStrictEqual(options, {job: fakeJob});
-          done();
+        getQueryResults: (options: {}, callback: Function) => {
+          callback(null, FAKE_ROWS, FAKE_RESPONSE);
         },
       };
 
       bq.createQueryJob = (query: {}, callback: Function) => {
         callback(null, fakeJob, FAKE_RESPONSE);
       };
 
-      bq.query(QUERY_STRING, fakeOptions, assert.ifError);
+      bq.query(MODEL_QUERY_STRING, (err: Error, rows: {}, resp: {}) => {
+        assert.ifError(err);
+        assert.strictEqual(rows, FAKE_ROWS);
+        assert.strictEqual(resp, FAKE_RESPONSE);
+        done();
+      });
     });
   });
 
@@ -1844,5 +1858,25 @@ describe('BigQuery', () => {
       };
       bq.queryAsStream_(query, done);
     });
+
+    it('should call query correctly with a job', done => {
+      const FAKE_ROWS = [{}, {}, {}];
+      const FAKE_RESPONSE = {};
+      const fakeJob = {
+        getQueryResults: (query: {}, callback: Function) => {
+          assert.strictEqual(query, query);
+          callback(null, FAKE_ROWS, FAKE_RESPONSE);
+        },
+      };
+
+      const query = {job: fakeJob};
+
+      bq.query = (query_: {}, options: {}, callback: Function) => {
+        assert.strictEqual(query_, query);
+        callback(); // done()
+      };
+
+      bq.queryAsStream_(query, done);
+    });
   });
 });