Add JS driver support for multiple results

brianc · brianc · commit 28b330c88eba · 2017-07-16T16:25:35.000-05:00
diff --git a/lib/query.js b/lib/query.js
@@ -29,10 +29,14 @@ var Query = function (config, values, callback) {
   // use unique portal name each time
   this.portal = config.portal || ''
   this.callback = config.callback
+  this._rowMode = config.rowMode
   if (process.domain && config.callback) {
     this.callback = process.domain.bind(config.callback)
   }
-  this._result = new Result(config.rowMode, config.types)
+  this._result = new Result(this._rowMode, this.types)
+
+  // potential for multiple results
+  this._results = this._result
   this.isPreparedStatement = false
   this._canceledDueToError = false
   this._promise = null
@@ -54,10 +58,24 @@ Query.prototype.requiresPreparation = function () {
   return this.values.length > 0
 }
 
+Query.prototype._checkForMultirow = function () {
+  // if we already have a result with a command property
+  // then we've already executed one query in a multi-statement simple query
+  // turn our results into an array of results
+  if (this._result.command) {
+    if (!Array.isArray(this._results)) {
+      this._results = [this._result]
+    }
+    this._result = new Result(this._rowMode, this.types)
+    this._results.push(this._result)
+  }
+}
+
 // associates row metadata from the supplied
 // message with this query object
 // metadata used when parsing row results
 Query.prototype.handleRowDescription = function (msg) {
+  this._checkForMultirow()
   this._result.addFields(msg.fields)
   this._accumulateRows = this.callback || !this.listeners('row').length
 }
@@ -83,6 +101,7 @@ Query.prototype.handleDataRow = function (msg) {
 }
 
 Query.prototype.handleCommandComplete = function (msg, con) {
+  this._checkForMultirow()
   this._result.addCommandComplete(msg)
   // need to sync after each command complete of a prepared statement
   if (this.isPreparedStatement) {
@@ -104,9 +123,9 @@ Query.prototype.handleReadyForQuery = function (con) {
     return this.handleError(this._canceledDueToError, con)
   }
   if (this.callback) {
-    this.callback(null, this._result)
+    this.callback(null, this._results)
   }
-  this.emit('end', this._result)
+  this.emit('end', this._results)
 }
 
 Query.prototype.handleError = function (err, connection) {
diff --git a/test/integration/client/multiple-results-tests.js b/test/integration/client/multiple-results-tests.js
@@ -0,0 +1,69 @@
+'use strict'
+const assert = require('assert')
+const co = require('co')
+
+const helper = require('./test-helper')
+
+const suite = new helper.Suite('multiple result sets')
+
+suite.test('two select results work', co.wrap(function * () {
+  const client = new helper.Client()
+  yield client.connect()
+
+  const results = yield client.query(`SELECT 'foo'::text as name; SELECT 'bar'::text as baz`)
+  assert(Array.isArray(results))
+
+  assert.equal(results[0].fields[0].name, 'name')
+  assert.deepEqual(results[0].rows, [{ name: 'foo' }])
+
+  assert.equal(results[1].fields[0].name, 'baz')
+  assert.deepEqual(results[1].rows, [{ baz: 'bar' }])
+
+  return client.end()
+}))
+
+suite.test('multiple selects work', co.wrap(function * () {
+  const client = new helper.Client()
+  yield client.connect()
+
+  const text = `
+  SELECT * FROM generate_series(2, 4) as foo;
+  SELECT * FROM generate_series(8, 10) as bar;
+  SELECT * FROM generate_series(20, 22) as baz;
+  `
+
+  const results = yield client.query(text)
+  assert(Array.isArray(results))
+
+  assert.equal(results[0].fields[0].name, 'foo')
+  assert.deepEqual(results[0].rows, [{ foo: 2 }, { foo: 3 }, { foo: 4 }])
+
+  assert.equal(results[1].fields[0].name, 'bar')
+  assert.deepEqual(results[1].rows, [{ bar: 8 }, { bar: 9 }, { bar: 10 }])
+
+  assert.equal(results[2].fields[0].name, 'baz')
+  assert.deepEqual(results[2].rows, [{ baz: 20 }, { baz: 21 }, { baz: 22 }])
+
+  assert.equal(results.length, 3)
+
+  return client.end()
+}))
+
+suite.test('mixed queries and statements', co.wrap(function * () {
+  const client = new helper.Client()
+  yield client.connect()
+
+  const text = `
+  CREATE TEMP TABLE weather(type text);
+  INSERT INTO weather(type) VALUES ('rain');
+  SELECT * FROM weather;
+  `
+
+  const results = yield client.query(text)
+  assert(Array.isArray(results))
+  assert.equal(results[0].command, 'CREATE')
+  assert.equal(results[1].command, 'INSERT')
+  assert.equal(results[2].command, 'SELECT')
+
+  return client.end()
+}))
