wip

Author: avallete

package.json

@@ -36,5 +36,6 @@
     "arrowParens": "always",
     "trailingComma": "es5",
     "singleQuote": true
-  }
+  },
+  "packageManager": "yarn@1.22.22+sha512.a6b2f7906b721bba3d67d4aff083df04dad64c399707841b7acf00f6b133b7ac24255f2652fa22ae3534329dc6180534e98d17432037ff6fd140556e2bb3137e"
 }

packages/pg/README.md

@@ -21,12 +21,45 @@ $ npm install pg
 - Pure JavaScript client and native libpq bindings share _the same API_
 - Connection pooling
 - Extensible JS ↔ PostgreSQL data-type coercion
+- Memory safety with configurable result size limits
 - Supported PostgreSQL features
   - Parameterized queries
   - Named statements with query plan caching
   - Async notifications with `LISTEN/NOTIFY`
   - Bulk import & export with `COPY TO/COPY FROM`
 
+### Memory Safety with Result Size Limits
+
+To prevent out-of-memory errors when dealing with unexpectedly large query results, you can set a maximum result size:
+
+```js
+const { Client, Pool } = require('pg')
+
+// For a single client
+const client = new Client({
+  // other connection options
+  maxResultSize: 50 * 1024 * 1024 // 50MB limit
+})
+
+// Or with a pool
+const pool = new Pool({
+  // other connection options
+  maxResultSize: 10 * 1024 * 1024 // 10MB limit
+})
+
+// If a query result exceeds the limit, it will emit an error with code 'RESULT_SIZE_EXCEEDED'
+client.query('SELECT * FROM large_table').catch(err => {
+  if (err.code === 'RESULT_SIZE_EXCEEDED') {
+    console.error(`Query result exceeded size limit of ${err.maxResultSize} bytes`)
+    // Handle gracefully - perhaps use a cursor or add a LIMIT clause
+  } else {
+    // Handle other errors
+  }
+})
+```
+
+For large datasets, consider using [pg-cursor](https://github.com/brianc/node-postgres/tree/master/packages/pg-cursor) to process rows in batches.
+
 ### Extras
 
 node-postgres is by design pretty light on abstractions. These are some handy modules we've been using over the years to complete the picture.

packages/pg/lib/client.js

@@ -52,6 +52,7 @@ class Client extends EventEmitter {
         keepAlive: c.keepAlive || false,
         keepAliveInitialDelayMillis: c.keepAliveInitialDelayMillis || 0,
         encoding: this.connectionParameters.client_encoding || 'utf8',
+        maxResultSize: c.maxResultSize
       })
     this.queryQueue = []
     this.binary = c.binary || defaults.binary
@@ -243,7 +244,7 @@ class Client extends EventEmitter {
     }
   }
 
-  _handleAuthCleartextPassword(msg) {
+  _handleAuthCleartextPassword(_msg) {
     this._checkPgPass(() => {
       this.connection.password(this.password)
     })
@@ -299,7 +300,7 @@ class Client extends EventEmitter {
     this.secretKey = msg.secretKey
   }
 
-  _handleReadyForQuery(msg) {
+  _handleReadyForQuery(_msg) {
     if (this._connecting) {
       this._connecting = false
       this._connected = true
@@ -376,12 +377,12 @@ class Client extends EventEmitter {
     this.activeQuery.handleDataRow(msg)
   }
 
-  _handlePortalSuspended(msg) {
+  _handlePortalSuspended(_msg) {
     // delegate portalSuspended to active query
     this.activeQuery.handlePortalSuspended(this.connection)
   }
 
-  _handleEmptyQuery(msg) {
+  _handleEmptyQuery(_msg) {
     // delegate emptyQuery to active query
     this.activeQuery.handleEmptyQuery(this.connection)
   }
@@ -410,7 +411,7 @@ class Client extends EventEmitter {
     }
   }
 
-  _handleCopyInResponse(msg) {
+  _handleCopyInResponse(_msg) {
     this.activeQuery.handleCopyInResponse(this.connection)
   }
 

packages/pg/lib/connection.js

@@ -27,6 +27,8 @@ class Connection extends EventEmitter {
     this.ssl = config.ssl || false
     this._ending = false
     this._emitMessage = false
+    this._maxResultSize = config.maxResultSize
+    this._currentResultSize = 0
     var self = this
     this.on('newListener', function (eventName) {
       if (eventName === 'message') {
@@ -108,6 +110,16 @@ class Connection extends EventEmitter {
   }
 
   attachListeners(stream) {
+    // Use the appropriate implementation based on whether maxResultSize is enabled
+    if (this._maxResultSize && this._maxResultSize > 0) {
+      this._attachListenersWithSizeLimit(stream)
+    } else {
+      this._attachListenersStandard(stream)
+    }
+  }
+
+  // Original implementation with no overhead
+  _attachListenersStandard(stream) {
     parse(stream, (msg) => {
       var eventName = msg.name === 'error' ? 'errorMessage' : msg.name
       if (this._emitMessage) {
@@ -117,6 +129,59 @@ class Connection extends EventEmitter {
     })
   }
 
+  // Implementation with size limiting logic
+  _attachListenersWithSizeLimit(stream) {
+    parse(stream, (msg) => {
+      var eventName = msg.name === 'error' ? 'errorMessage' : msg.name
+      
+      // Only track data row messages for result size
+      if (msg.name === 'dataRow') {
+        // Approximate size by using message length
+        const msgSize = msg.length || this._getApproximateMessageSize(msg)
+        this._currentResultSize += msgSize
+        
+        // Check if we've exceeded the max result size
+        if (this._currentResultSize > this._maxResultSize) {
+          const error = new Error('Query result size exceeded the configured limit')
+          error.code = 'RESULT_SIZE_EXCEEDED'
+          error.resultSize = this._currentResultSize
+          error.maxResultSize = this._maxResultSize
+          this.emit('error', error)
+          this.end() // Terminate the connection
+          return
+        }
+      }
+      
+      // Reset counter on query completion
+      if (msg.name === 'readyForQuery') {
+        this._currentResultSize = 0
+      }
+      
+      if (this._emitMessage) {
+        this.emit('message', msg)
+      }
+      this.emit(eventName, msg)
+    })
+  }
+
+  // Helper method to approximate message size when length is not available
+  _getApproximateMessageSize(msg) {
+    let size = 0
+    if (msg.fields) {
+      // Sum up the sizes of field values
+      msg.fields.forEach(field => {
+        if (field && typeof field === 'string') {
+          size += field.length;
+        } else if (field && typeof field === 'object') {
+          size += JSON.stringify(field).length;
+        } else if (field !== null && field !== undefined) {
+          size += String(field).length;
+        }
+      });
+    }
+    return size > 0 ? size : 1024; // Default to 1KB if we can't determine size
+  }
+
   requestSsl() {
     this.stream.write(serialize.requestSsl())
   }

packages/pg/lib/defaults.js

@@ -70,6 +70,8 @@ module.exports = {
   keepalives: 1,
 
   keepalives_idle: 0,
+  // maxResultSize limit of a request before erroring out
+  maxResultSize: undefined,
 }
 
 var pgTypes = require('pg-types')

packages/pg/test/integration/client/max-result-size-tests.js

@@ -0,0 +1,98 @@
+'use strict'
+var helper = require('./test-helper')
+var assert = require('assert')
+const { Client } = helper.pg
+const suite = new helper.Suite()
+
+// Just verify the test infrastructure works
+suite.test('sanity check', function (done) {
+  var client = new Client(helper.args)
+  client.connect(assert.success(function () {
+    client.query('SELECT 1 as num', assert.success(function(result) {
+      assert.equal(result.rows.length, 1)
+      client.end(done)
+    }))
+  }))
+})
+
+// Basic test to check if the _maxResultSize property is passed to Connection
+suite.test('client passes maxResultSize to connection', function (done) {
+  var client = new Client({
+    ...helper.args,
+    maxResultSize: 1024 * 1024 // 1MB limit
+  })
+  
+  client.connect(assert.success(function () {
+    assert.equal(client.connection._maxResultSize, 1024 * 1024, 
+                 'maxResultSize should be passed to the connection')
+    client.end(done)
+  }))
+})
+
+// Check if the correct attachListeners method is being used based on maxResultSize
+suite.test('connection uses correct listener implementation', function (done) {
+  var client = new Client({
+    ...helper.args,
+    maxResultSize: 1024 * 1024 // 1MB limit
+  })
+  
+  client.connect(assert.success(function () {
+    // Just a simple check to see if our methods exist on the connection object
+    assert(typeof client.connection._attachListenersStandard === 'function',
+           'Standard listener method should exist')
+    assert(typeof client.connection._attachListenersWithSizeLimit === 'function',
+           'Size-limiting listener method should exist')
+    client.end(done)
+  }))
+})
+
+// Test that small result sets complete successfully with maxResultSize set
+suite.test('small result with maxResultSize', function (done) {
+  var client = new Client({
+    ...helper.args,
+    maxResultSize: 1024 * 1024 // 1MB limit
+  })
+  
+  client.connect(assert.success(function () {
+    client.query('SELECT generate_series(1, 10) as num', assert.success(function(result) {
+      assert.equal(result.rows.length, 10)
+      client.end(done)
+    }))
+  }))
+})
+
+// Test for large result size
+// Using a separate test to avoid issue with callback being called twice
+suite.testAsync('large result triggers error', async () => {
+  const client = new Client({
+    ...helper.args,
+    maxResultSize: 500 // Very small limit
+  })
+
+  // Setup error handler
+  const errorPromise = new Promise(resolve => {
+    client.on('error', err => {
+      assert.equal(err.message, 'Query result size exceeded the configured limit')
+      assert.equal(err.code, 'RESULT_SIZE_EXCEEDED')
+      resolve()
+    })
+  })
+
+  await client.connect()
+  
+  // Start the query but don't await it (it will error)
+  const queryPromise = client.query('SELECT repeat(\'x\', 1000) as data FROM generate_series(1, 100)')
+    .catch(err => {
+      // We expect this to error out, silence the rejection
+      return null
+    })
+  
+  // Wait for error event
+  await errorPromise
+  
+  // Make sure the query is done before we end
+  await queryPromise
+  
+  // Clean up
+  await client.end().catch(() => {}) // Ignore errors during cleanup
+})