diff --git a/packages/pg/lib/connection.js b/packages/pg/lib/connection.js index fe04efb6b..86724c5c5 100644 --- a/packages/pg/lib/connection.js +++ b/packages/pg/lib/connection.js @@ -14,7 +14,12 @@ class Connection extends EventEmitter { constructor(config) { super() config = config || {} + this.stream = config.stream || new net.Socket() + if (typeof this.stream === 'function') { + this.stream = this.stream(config) + } + this._keepAlive = config.keepAlive this._keepAliveInitialDelayMillis = config.keepAliveInitialDelayMillis this.lastBuffer = false diff --git a/packages/pg/test/unit/connection/startup-tests.js b/packages/pg/test/unit/connection/startup-tests.js index e2eb6ee99..d5d30d5de 100644 --- a/packages/pg/test/unit/connection/startup-tests.js +++ b/packages/pg/test/unit/connection/startup-tests.js @@ -7,6 +7,18 @@ test('connection can take existing stream', function () { assert.equal(con.stream, stream) }) +test('connection can take stream factory method', function () { + var stream = new MemoryStream() + var connectionOpts = {} + var makeStream = function (opts) { + assert.equal(connectionOpts, opts) + return stream + } + connectionOpts.stream = makeStream + var con = new Connection(connectionOpts) + assert.equal(con.stream, stream) +}) + test('using any stream', function () { var makeStream = function () { var stream = new MemoryStream()