chore: add tests for native

Author: avallete

packages/pg/lib/native/client.js

@@ -26,6 +26,9 @@ var Client = (module.exports = function (config) {
     types: this._types,
   })
 
+  // Store maxResultSize configuration
+  this._maxResultSize = config.maxResultSize
+
   this._queryQueue = []
   this._ending = false
   this._connecting = false
@@ -100,6 +103,9 @@ Client.prototype._connect = function (cb) {
       // set internal states to connected
       self._connected = true
 
+      // Add a reference to the client for error bubbling
+      self.native.connection = self;
+
       // handle connection errors from the native layer
       self.native.on('error', function (err) {
         self._queryable = false
@@ -196,7 +202,7 @@ Client.prototype.query = function (config, values, callback) {
 
       // we already returned an error,
       // just do nothing if query completes
-      query.callback = () => {}
+      query.callback = () => { }
 
       // Remove from queue
       var index = this._queryQueue.indexOf(query)
@@ -289,14 +295,14 @@ Client.prototype._pulseQueryQueue = function (initialConnection) {
 // attempt to cancel an in-progress query
 Client.prototype.cancel = function (query) {
   if (this._activeQuery === query) {
-    this.native.cancel(function () {})
+    this.native.cancel(function () { })
   } else if (this._queryQueue.indexOf(query) !== -1) {
     this._queryQueue.splice(this._queryQueue.indexOf(query), 1)
   }
 }
 
-Client.prototype.ref = function () {}
-Client.prototype.unref = function () {}
+Client.prototype.ref = function () { }
+Client.prototype.unref = function () { }
 
 Client.prototype.setTypeParser = function (oid, format, parseFn) {
   return this._types.setTypeParser(oid, format, parseFn)

packages/pg/lib/native/query.js

@@ -56,6 +56,17 @@ NativeQuery.prototype.handleError = function (err) {
       err[normalizedFieldName] = fields[key]
     }
   }
+
+  // For maxResultSize exceeded errors, make sure we emit the error to the client too
+  if (err.code === 'RESULT_SIZE_EXCEEDED') {
+    if (this.native && this.native.connection) {
+      // Need to emit the error on the client/connection level too
+      process.nextTick(() => {
+        this.native.connection.emit('error', err);
+      });
+    }
+  }
+
   if (this.callback) {
     this.callback(err)
   } else {
@@ -89,6 +100,9 @@ NativeQuery.prototype.submit = function (client) {
   this.native = client.native
   client.native.arrayMode = this._arrayMode
 
+  // Get the maxResultSize from the client if it's set
+  this._maxResultSize = client._maxResultSize
+
   var after = function (err, rows, results) {
     client.native.arrayMode = false
     setImmediate(function () {
@@ -100,6 +114,30 @@ NativeQuery.prototype.submit = function (client) {
       return self.handleError(err)
     }
 
+    // Check the result size if maxResultSize is configured
+    if (self._maxResultSize) {
+      // Calculate result size (rough approximation)
+      let resultSize = 0;
+
+      // For multiple result sets
+      if (results.length > 1) {
+        for (let i = 0; i < rows.length; i++) {
+          resultSize += self._calculateResultSize(rows[i]);
+        }
+      } else if (rows.length > 0) {
+        resultSize = self._calculateResultSize(rows);
+      }
+
+      // If the size limit is exceeded, generate an error
+      if (resultSize > self._maxResultSize) {
+        const error = new Error('Query result size exceeded the configured limit');
+        error.code = 'RESULT_SIZE_EXCEEDED';
+        error.resultSize = resultSize;
+        error.maxResultSize = self._maxResultSize;
+        return self.handleError(error);
+      }
+    }
+
     // emit row events for each row in the result
     if (self._emitRowEvents) {
       if (results.length > 1) {
@@ -166,3 +204,59 @@ NativeQuery.prototype.submit = function (client) {
     client.native.query(this.text, after)
   }
 }
+
+// Helper method to estimate the size of a result set
+NativeQuery.prototype._calculateResultSize = function (rows) {
+  let size = 0;
+
+  // For empty results, return 0
+  if (!rows || rows.length === 0) {
+    return 0;
+  }
+
+  // For array mode, calculate differently
+  if (this._arrayMode) {
+    // Just use a rough approximation based on number of rows
+    return rows.length * 100;
+  }
+
+  // For each row, approximate its size
+  for (let i = 0; i < rows.length; i++) {
+    const row = rows[i];
+
+    // Add base row size
+    size += 24; // Overhead per row
+
+    // Add size of each column
+    for (const key in row) {
+      if (Object.prototype.hasOwnProperty.call(row, key)) {
+        const value = row[key];
+
+        // Add key size
+        size += key.length * 2; // Assume 2 bytes per character
+
+        // Add value size based on type
+        if (value === null || value === undefined) {
+          size += 8;
+        } else if (typeof value === 'string') {
+          size += value.length * 2; // Assume 2 bytes per character
+        } else if (typeof value === 'number') {
+          size += 8;
+        } else if (typeof value === 'boolean') {
+          size += 4;
+        } else if (value instanceof Date) {
+          size += 8;
+        } else if (Buffer.isBuffer(value)) {
+          size += value.length;
+        } else if (Array.isArray(value)) {
+          size += 16 + value.length * 8;
+        } else {
+          // For objects, use a rough estimate
+          size += 32 + JSON.stringify(value).length * 2;
+        }
+      }
+    }
+  }
+
+  return size;
+}

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

@@ -10,18 +10,19 @@ process.on('unhandledRejection', function (e) {
 
 const suite = new helper.Suite()
 
-// Test that result size exceeding maxResultSize properly raises an error
 suite.test('maxResultSize limit triggers error', (cb) => {
+  // Check if we're running with the native client
+  const isNative = helper.args.native
+  console.log(isNative ? 'Testing with native client' : 'Testing with JavaScript client')
+
   // Create a pool with a very small result size limit
   const pool = new pg.Pool({
     maxResultSize: 100, // Very small limit (100 bytes)
     ...helper.args,
   })
 
-  // Track if we've seen the size exceeded error
   let sizeExceededErrorSeen = false
 
-  // Set up error handler on pool
   pool.on('error', (err) => {
     console.log('Pool error:', err.message, err.code)
   })
@@ -33,21 +34,20 @@ suite.test('maxResultSize limit triggers error', (cb) => {
       client.on('error', (err) => {
         console.log('Client error event:', err.message, err.code)
 
-        // If we get the expected error, mark it
-        if (err.message === 'Query result size exceeded the configured limit') {
-          assert.equal(err.code, 'RESULT_SIZE_EXCEEDED', 'Error should have RESULT_SIZE_EXCEEDED code')
+        // If we get any size exceeded error, mark it
+        if (err.code === 'RESULT_SIZE_EXCEEDED' ||
+          err.message === 'Query result size exceeded the configured limit') {
           sizeExceededErrorSeen = true
         }
       })
 
-      // Create a temp table
       return client
         .query('CREATE TEMP TABLE large_result_test(id SERIAL, data TEXT)')
         .then(() => {
-          // Insert data that will exceed the size limit when selected
+          // Insert rows that will exceed the size limit when queried
           const insertPromises = []
           for (let i = 0; i < 20; i++) {
-            // Each row will have 50 bytes of data
+            // Each row will have enough data to eventually exceed our limit
             const data = 'x'.repeat(50)
             insertPromises.push(client.query('INSERT INTO large_result_test(data) VALUES($1)', [data]))
           }
@@ -56,7 +56,6 @@ suite.test('maxResultSize limit triggers error', (cb) => {
         .then(() => {
           console.log('Running query that should exceed size limit...')
 
-          // This query should fail due to exceeding size limit
           return client
             .query('SELECT * FROM large_result_test')
             .then(() => {
@@ -65,15 +64,14 @@ suite.test('maxResultSize limit triggers error', (cb) => {
             .catch((err) => {
               console.log('Query error caught:', err.message, err.code)
 
-              // The error should have the correct code
+              // Both implementations should throw an error with this code
               assert.equal(err.code, 'RESULT_SIZE_EXCEEDED', 'Error should have RESULT_SIZE_EXCEEDED code')
 
-              // Give a little time for error events to be processed
+              // Give time for error events to propagate
               return new Promise((resolve) => setTimeout(resolve, 100)).then(() => {
-                // Verify we saw the expected error event
+                // Verify we saw the error event
                 assert(sizeExceededErrorSeen, 'Should have seen the size exceeded error event')
 
-                // Attempt cleanup but don't fail if it errors
                 return client.query('DROP TABLE IF EXISTS large_result_test').catch(() => {
                   /* ignore cleanup errors */
                 })
@@ -85,16 +83,17 @@ suite.test('maxResultSize limit triggers error', (cb) => {
           pool.end(cb)
         })
         .catch((err) => {
+          console.error('Test error:', err.message)
           client.release()
           pool.end(() => cb(err))
         })
     })
     .catch((err) => {
+      console.error('Connection error:', err.message)
       pool.end(() => cb(err))
     })
 })
 
-// Test that results under the maxResultSize limit work normally
 suite.test('results under maxResultSize limit work correctly', (cb) => {
   // Create a pool with a reasonably large limit
   const pool = new pg.Pool({
@@ -105,21 +104,16 @@ suite.test('results under maxResultSize limit work correctly', (cb) => {
   pool
     .connect()
     .then((client) => {
-      // Create a temp table
       return client
         .query('CREATE TEMP TABLE small_result_test(id SERIAL, data TEXT)')
         .then(() => {
-          // Insert a small amount of data
           return client.query('INSERT INTO small_result_test(data) VALUES($1)', ['small_data'])
         })
         .then(() => {
-          // This query should succeed
           return client.query('SELECT * FROM small_result_test').then((result) => {
-            // Verify the result
             assert.equal(result.rows.length, 1, 'Should get 1 row')
             assert.equal(result.rows[0].data, 'small_data', 'Data should match')
 
-            // Clean up
             return client.query('DROP TABLE small_result_test')
           })
         })