Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 0e786d2

Browse files
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+
}
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)

AltStyle によって変換されたページ (->オリジナル) /