Skip to content

Commit 0e786d2

Browse files
avalleteavallete
committed
feat: add maxResultSize limit
1 parent 477f812 commit 0e786d2

File tree

6 files changed

+282
-0
lines changed

6 files changed

+282
-0
lines changed

‎packages/pg/lib/client.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ class Client extends EventEmitter {
5252
keepAlive: c.keepAlive || false,
5353
keepAliveInitialDelayMillis: c.keepAliveInitialDelayMillis || 0,
5454
encoding: this.connectionParameters.client_encoding || 'utf8',
55+
maxResultSize: c.maxResultSize,
5556
})
5657
this.queryQueue = []
5758
this.binary = c.binary || defaults.binary

‎packages/pg/lib/connection.js‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ class Connection extends EventEmitter {
2727
this.ssl = config.ssl || false
2828
this._ending = false
2929
this._emitMessage = false
30+
this._maxResultSize = config.maxResultSize
31+
this._currentResultSize = 0
3032
var self = this
3133
this.on('newListener', function (eventName) {
3234
if (eventName === 'message') {
@@ -108,6 +110,17 @@ class Connection extends EventEmitter {
108110
}
109111

110112
attachListeners(stream) {
113+
var self = this
114+
// Use the appropriate implementation based on whether maxResultSize is enabled
115+
if (self._maxResultSize && self._maxResultSize > 0) {
116+
this._attachListenersWithSizeLimit(stream)
117+
} else {
118+
this._attachListenersStandard(stream)
119+
}
120+
}
121+
122+
// Original implementation with no overhead
123+
_attachListenersStandard(stream) {
111124
parse(stream, (msg) => {
112125
var eventName = msg.name === 'error' ? 'errorMessage' : msg.name
113126
if (this._emitMessage) {
@@ -117,6 +130,41 @@ class Connection extends EventEmitter {
117130
})
118131
}
119132

133+
// Implementation with size limiting logic
134+
_attachListenersWithSizeLimit(stream) {
135+
parse(stream, (msg) => {
136+
var eventName = msg.name === 'error' ? 'errorMessage' : msg.name
137+
138+
// Only track data row messages for result size
139+
if (msg.name === 'dataRow') {
140+
// Approximate size by using message length
141+
const msgSize = msg.length || 1024 // Default to 1KB if we don't have lenght info
142+
this._currentResultSize += msgSize
143+
144+
// Check if we've exceeded the max result size
145+
if (this._currentResultSize > this._maxResultSize) {
146+
const error = new Error('Query result size exceeded the configured limit')
147+
error.code = 'RESULT_SIZE_EXCEEDED'
148+
error.resultSize = this._currentResultSize
149+
error.maxResultSize = this._maxResultSize
150+
this.emit('errorMessage', error)
151+
this.end() // Terminate the connection
152+
return
153+
}
154+
}
155+
156+
// Reset counter on query completion
157+
if (msg.name === 'readyForQuery') {
158+
this._currentResultSize = 0
159+
}
160+
161+
if (this._emitMessage) {
162+
this.emit('message', msg)
163+
}
164+
this.emit(eventName, msg)
165+
})
166+
}
167+
120168
requestSsl() {
121169
this.stream.write(serialize.requestSsl())
122170
}

‎packages/pg/lib/defaults.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ module.exports = {
7070
keepalives: 1,
7171

7272
keepalives_idle: 0,
73+
// maxResultSize limit of a request before erroring out
74+
maxResultSize: undefined,
7375
}
7476

7577
var pgTypes = require('pg-types')

‎packages/pg/lib/native/client.js‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ var Client = (module.exports = function (config) {
2626
types: this._types,
2727
})
2828

29+
// Store maxResultSize configuration
30+
this._maxResultSize = config.maxResultSize
31+
2932
this._queryQueue = []
3033
this._ending = false
3134
this._connecting = false
@@ -100,6 +103,9 @@ Client.prototype._connect = function (cb) {
100103
// set internal states to connected
101104
self._connected = true
102105

106+
// Add a reference to the client for error bubbling
107+
self.native.connection = self
108+
103109
// handle connection errors from the native layer
104110
self.native.on('error', function (err) {
105111
self._queryable = false

‎packages/pg/lib/native/query.js‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,17 @@ NativeQuery.prototype.handleError = function (err) {
5656
err[normalizedFieldName] = fields[key]
5757
}
5858
}
59+
60+
// For maxResultSize exceeded errors, make sure we emit the error to the client too
61+
if (err.code === 'RESULT_SIZE_EXCEEDED') {
62+
if (this.native && this.native.connection) {
63+
// Need to emit the error on the client/connection level too
64+
process.nextTick(() => {
65+
this.native.connection.emit('error', err)
66+
})
67+
}
68+
}
69+
5970
if (this.callback) {
6071
this.callback(err)
6172
} else {
@@ -89,6 +100,9 @@ NativeQuery.prototype.submit = function (client) {
89100
this.native = client.native
90101
client.native.arrayMode = this._arrayMode
91102

103+
// Get the maxResultSize from the client if it's set
104+
this._maxResultSize = client._maxResultSize
105+
92106
var after = function (err, rows, results) {
93107
client.native.arrayMode = false
94108
setImmediate(function () {
@@ -100,6 +114,30 @@ NativeQuery.prototype.submit = function (client) {
100114
return self.handleError(err)
101115
}
102116

117+
// Check the result size if maxResultSize is configured
118+
if (self._maxResultSize) {
119+
// Calculate result size (rough approximation)
120+
let resultSize = 0
121+
122+
// For multiple result sets
123+
if (results.length > 1) {
124+
for (let i = 0; i < rows.length; i++) {
125+
resultSize += self._calculateResultSize(rows[i])
126+
}
127+
} else if (rows.length > 0) {
128+
resultSize = self._calculateResultSize(rows)
129+
}
130+
131+
// If the size limit is exceeded, generate an error
132+
if (resultSize > self._maxResultSize) {
133+
const error = new Error('Query result size exceeded the configured limit')
134+
error.code = 'RESULT_SIZE_EXCEEDED'
135+
error.resultSize = resultSize
136+
error.maxResultSize = self._maxResultSize
137+
return self.handleError(error)
138+
}
139+
}
140+
103141
// emit row events for each row in the result
104142
if (self._emitRowEvents) {
105143
if (results.length > 1) {
@@ -166,3 +204,59 @@ NativeQuery.prototype.submit = function (client) {
166204
client.native.query(this.text, after)
167205
}
168206
}
207+
208+
// Helper method to estimate the size of a result set
209+
NativeQuery.prototype._calculateResultSize = function (rows) {
210+
let size = 0
211+
212+
// For empty results, return 0
213+
if (!rows || rows.length === 0) {
214+
return 0
215+
}
216+
217+
// For array mode, calculate differently
218+
if (this._arrayMode) {
219+
// Just use a rough approximation based on number of rows
220+
return rows.length * 100
221+
}
222+
223+
// For each row, approximate its size
224+
for (let i = 0; i < rows.length; i++) {
225+
const row = rows[i]
226+
227+
// Add base row size
228+
size += 24 // Overhead per row
229+
230+
// Add size of each column
231+
for (const key in row) {
232+
if (Object.prototype.hasOwnProperty.call(row, key)) {
233+
const value = row[key]
234+
235+
// Add key size
236+
size += key.length * 2 // Assume 2 bytes per character
237+
238+
// Add value size based on type
239+
if (value === null || value === undefined) {
240+
size += 8
241+
} else if (typeof value === 'string') {
242+
size += value.length * 2 // Assume 2 bytes per character
243+
} else if (typeof value === 'number') {
244+
size += 8
245+
} else if (typeof value === 'boolean') {
246+
size += 4
247+
} else if (value instanceof Date) {
248+
size += 8
249+
} else if (Buffer.isBuffer(value)) {
250+
size += value.length
251+
} else if (Array.isArray(value)) {
252+
size += 16 + value.length * 8
253+
} else {
254+
// For objects, use a rough estimate
255+
size += 32 + JSON.stringify(value).length * 2
256+
}
257+
}
258+
}
259+
}
260+
261+
return size
262+
}

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
'use strict'
2+
const helper = require('../test-helper')
3+
const pg = helper.pg
4+
const assert = require('assert')
5+
6+
process.on('unhandledRejection', function (e) {
7+
console.error(e, e.stack)
8+
process.exit(1)
9+
})
10+
11+
const suite = new helper.Suite()
12+
13+
suite.test('maxResultSize limit triggers error', (cb) => {
14+
// Check if we're running with the native client
15+
const isNative = helper.args.native
16+
console.log(isNative ? 'Testing with native client' : 'Testing with JavaScript client')
17+
18+
// Create a pool with a very small result size limit
19+
const pool = new pg.Pool({
20+
maxResultSize: 100, // Very small limit (100 bytes)
21+
...helper.args,
22+
})
23+
24+
let sizeExceededErrorSeen = false
25+
26+
pool.on('error', (err) => {
27+
console.log('Pool error:', err.message, err.code)
28+
})
29+
30+
pool
31+
.connect()
32+
.then((client) => {
33+
// Set up client error listener for error events
34+
client.on('error', (err) => {
35+
console.log('Client error event:', err.message, err.code)
36+
37+
// If we get any size exceeded error, mark it
38+
if (err.code === 'RESULT_SIZE_EXCEEDED' || err.message === 'Query result size exceeded the configured limit') {
39+
sizeExceededErrorSeen = true
40+
}
41+
})
42+
43+
return client
44+
.query('CREATE TEMP TABLE large_result_test(id SERIAL, data TEXT)')
45+
.then(() => {
46+
// Insert rows that will exceed the size limit when queried
47+
const insertPromises = []
48+
for (let i = 0; i < 20; i++) {
49+
// Each row will have enough data to eventually exceed our limit
50+
const data = 'x'.repeat(50)
51+
insertPromises.push(client.query('INSERT INTO large_result_test(data) VALUES($1)', [data]))
52+
}
53+
return Promise.all(insertPromises)
54+
})
55+
.then(() => {
56+
console.log('Running query that should exceed size limit...')
57+
58+
return client
59+
.query('SELECT * FROM large_result_test')
60+
.then(() => {
61+
throw new Error('Query should have failed due to size limit')
62+
})
63+
.catch((err) => {
64+
console.log('Query error caught:', err.message, err.code)
65+
66+
// Both implementations should throw an error with this code
67+
assert.equal(err.code, 'RESULT_SIZE_EXCEEDED', 'Error should have RESULT_SIZE_EXCEEDED code')
68+
69+
// Give time for error events to propagate
70+
return new Promise((resolve) => setTimeout(resolve, 100)).then(() => {
71+
// Verify we saw the error event
72+
assert(sizeExceededErrorSeen, 'Should have seen the size exceeded error event')
73+
74+
return client.query('DROP TABLE IF EXISTS large_result_test').catch(() => {
75+
/* ignore cleanup errors */
76+
})
77+
})
78+
})
79+
})
80+
.then(() => {
81+
client.release()
82+
pool.end(cb)
83+
})
84+
.catch((err) => {
85+
console.error('Test error:', err.message)
86+
client.release()
87+
pool.end(() => cb(err))
88+
})
89+
})
90+
.catch((err) => {
91+
console.error('Connection error:', err.message)
92+
pool.end(() => cb(err))
93+
})
94+
})
95+
96+
suite.test('results under maxResultSize limit work correctly', (cb) => {
97+
// Create a pool with a reasonably large limit
98+
const pool = new pg.Pool({
99+
maxResultSize: 10 * 1024, // 10KB is plenty for small results
100+
...helper.args,
101+
})
102+
103+
pool
104+
.connect()
105+
.then((client) => {
106+
return client
107+
.query('CREATE TEMP TABLE small_result_test(id SERIAL, data TEXT)')
108+
.then(() => {
109+
return client.query('INSERT INTO small_result_test(data) VALUES($1)', ['small_data'])
110+
})
111+
.then(() => {
112+
return client.query('SELECT * FROM small_result_test').then((result) => {
113+
assert.equal(result.rows.length, 1, 'Should get 1 row')
114+
assert.equal(result.rows[0].data, 'small_data', 'Data should match')
115+
116+
return client.query('DROP TABLE small_result_test')
117+
})
118+
})
119+
.then(() => {
120+
client.release()
121+
pool.end(cb)
122+
})
123+
.catch((err) => {
124+
client.release()
125+
pool.end(() => cb(err))
126+
})
127+
})
128+
.catch((err) => {
129+
pool.end(() => cb(err))
130+
})
131+
})

0 commit comments

Comments
 (0)