Skip to content

Commit c61864b

Browse files
adrianprelipceanadrianprelipcean
committed
feat: Changed raw SQL calls to stored procedures.
1 parent a8be6c3 commit c61864b

File tree

9 files changed

+132
-88
lines changed

9 files changed

+132
-88
lines changed

‎source/migrations/2-add-db-functions-agg.sql

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
CREATE OR REPLACE FUNCTION rate_limit.agg_increment(key_ text, session_id_ uuid)
2+
RETURNS int AS
3+
$bd$
4+
INSERT INTO rate_limit.records_aggregated(key, session_id)
5+
VALUES ($1, $2)
6+
ON CONFLICT ON CONSTRAINT unique_session_key DO UPDATE
7+
SET count = records_aggregated.count + 1
8+
RETURNING count;
9+
$bd$
10+
LANGUAGE SQL;
11+
12+
CREATE OR REPLACE FUNCTION rate_limit.agg_decrement(key_ text, session_id_ uuid)
13+
RETURNS void AS
14+
$bd$
15+
UPDATE rate_limit.records_aggregated
16+
SET count = greatest(0, count-1)
17+
WHERE key = $1 and session_id = $2;
18+
$bd$
19+
LANGUAGE SQL;
20+
21+
CREATE OR REPLACE FUNCTION rate_limit.agg_reset_key(key_ text, session_id_ uuid)
22+
RETURNS void AS
23+
$bd$
24+
DELETE FROM rate_limit.records_aggregated
25+
WHERE key = $1 and session_id = $2
26+
$bd$
27+
LANGUAGE SQL;
28+
29+
CREATE OR REPLACE FUNCTION rate_limit.agg_reset_session(session_id_ uuid)
30+
RETURNS void AS
31+
$bd$
32+
DELETE FROM rate_limit.records_aggregated
33+
WHERE session_id = $1
34+
$bd$
35+
LANGUAGE SQL;

‎source/migrations/3-add-db-functions-ind.sql

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
CREATE OR REPLACE FUNCTION rate_limit.ind_increment(key_ text, session_id_ uuid)
2+
RETURNS int AS
3+
$bd$
4+
INSERT INTO rate_limit.individual_records(key, session_id) VALUES ($1, $2);
5+
6+
SELECT count(id) AS count FROM rate_limit.individual_records WHERE key = $1 AND session_id = $2;
7+
$bd$
8+
LANGUAGE SQL;
9+
10+
CREATE OR REPLACE FUNCTION rate_limit.ind_decrement(key_ text, session_id_ uuid)
11+
RETURNS void AS
12+
$bd$
13+
WITH
14+
rows_to_delete AS (
15+
SELECT id FROM rate_limit.individual_records
16+
WHERE key = $1 and session_id = $2 ORDER BY event_time LIMIT 1
17+
)
18+
DELETE FROM rate_limit.individual_records
19+
USING rows_to_delete WHERE individual_records.id = rows_to_delete.id
20+
$bd$
21+
LANGUAGE SQL;
22+
23+
CREATE OR REPLACE FUNCTION rate_limit.ind_reset_key(key_ text, session_id_ uuid)
24+
RETURNS void AS
25+
$bd$
26+
DELETE FROM rate_limit.individual_records
27+
WHERE key = $1 AND session_id = $2
28+
$bd$
29+
LANGUAGE SQL;
30+
31+
CREATE OR REPLACE FUNCTION rate_limit.ind_reset_session(session_id_ uuid)
32+
RETURNS void AS
33+
$bd$
34+
DELETE FROM rate_limit.individual_records
35+
WHERE session_id = $1
36+
$bd$
37+
LANGUAGE SQL;

‎source/migrations/4-add-db-functions-sessions.sql

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
CREATE OR REPLACE FUNCTION rate_limit.session_select(name_ text, type_ text)
2+
RETURNS TABLE(id uuid, name_ text, type_ text, expires_at timestamptz) AS
3+
$bd$
4+
SELECT id, name_, type_, expires_at
5+
FROM rate_limit.sessions
6+
WHERE name_ = $1 and type_ = $2
7+
LIMIT 1;
8+
$bd$
9+
LANGUAGE SQL;
10+
11+
CREATE OR REPLACE FUNCTION rate_limit.session_reset(name_ text, type_ text, expires_at_ timestamptz)
12+
RETURNS TABLE(id uuid, name_ text, type_ text) AS
13+
$bd$
14+
DELETE FROM rate_limit.sessions
15+
WHERE name_ = $1 and type_ = $2;
16+
17+
INSERT INTO rate_limit.sessions(name_, type_, expires_at)
18+
SELECT $1, $2, $3
19+
RETURNING id, name_, type_;
20+
$bd$
21+
LANGUAGE SQL;

‎source/stores/aggregated_ip/store_aggregated_ip.ts

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,7 @@ class PostgresStoreAggregatedIP implements Store {
9191
*/
9292
async increment(key: string): Promise<ClientRateLimitInfo> {
9393
let recordInsertGetRecordsQuery = `
94-
INSERT INTO rate_limit.records_aggregated(key, session_id) VALUES ($1, $2)
95-
ON CONFLICT ON CONSTRAINT unique_session_key DO UPDATE
96-
SET count = records_aggregated.count + 1
97-
RETURNING count
94+
SELECT agg_increment as count FROM rate_limit.agg_increment($1, $2);
9895
`
9996
if (!isSessionValid(this.session)) {
10097
this.session = await getSession(
@@ -132,9 +129,7 @@ class PostgresStoreAggregatedIP implements Store {
132129
*/
133130
async decrement(key: string): Promise<void> {
134131
let decrementQuery = `
135-
UPDATE rate_limit.records_aggregated
136-
SET count = greatest(0, count-1)
137-
WHERE key = $1 and session_id = $2
132+
SELECT * FROM rate_limit.agg_decrement($1, $2);
138133
`
139134

140135
try {
@@ -154,8 +149,7 @@ class PostgresStoreAggregatedIP implements Store {
154149
*/
155150
async resetKey(key: string): Promise<void> {
156151
let resetQuery = `
157-
DELETE FROM rate_limit.records_aggregated
158-
WHERE key = $1 and session_id = $2
152+
SELECT * FROM rate_limit.agg_reset_key($1, $2)
159153
`
160154

161155
try {
@@ -175,7 +169,7 @@ class PostgresStoreAggregatedIP implements Store {
175169
*/
176170
async resetAll(): Promise<void> {
177171
let resetAllQuery = `
178-
DELETE FROM rate_limit.records_aggregated WHERE session_id = $1
172+
SELECT * FROM rate_limit.agg_reset_session($1);
179173
`
180174

181175
try {

‎source/stores/individual_ip/store_individual_ip.ts

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,7 @@ class PostgresStoreIndividualIP implements Store {
9191
*/
9292
async increment(key: string): Promise<ClientRateLimitInfo> {
9393
let recordInsertQuery =
94-
'INSERT INTO rate_limit.individual_records(key, session_id) VALUES ($1, $2)'
95-
let numberOfRecordsQuery =
96-
'SELECT count(id) AS count FROM rate_limit.individual_records WHERE key = $1 AND session_id = $2'
94+
'SELECT ind_increment as count FROM rate_limit.ind_increment($1, $2)'
9795
if (!isSessionValid(this.session)) {
9896
this.session = await getSession(
9997
this.prefix,
@@ -104,8 +102,7 @@ class PostgresStoreIndividualIP implements Store {
104102
}
105103

106104
try {
107-
await this.pool.query(recordInsertQuery, [key, this.session.id])
108-
let result = await this.pool.query(numberOfRecordsQuery, [
105+
let result = await this.pool.query(recordInsertQuery, [
109106
key,
110107
this.session.id,
111108
])
@@ -131,13 +128,7 @@ class PostgresStoreIndividualIP implements Store {
131128
*/
132129
async decrement(key: string): Promise<void> {
133130
let decrementQuery = `
134-
WITH
135-
rows_to_delete AS (
136-
SELECT id FROM rate_limit.individual_records
137-
WHERE key = $1 and session_id = $2 ORDER BY event_time LIMIT 1
138-
)
139-
DELETE FROM rate_limit.individual_records
140-
USING rows_to_delete WHERE individual_records.id = rows_to_delete.id
131+
SELECT * FROM rate_limit.ind_decrement($1, $2);
141132
`
142133

143134
try {
@@ -157,8 +148,7 @@ class PostgresStoreIndividualIP implements Store {
157148
*/
158149
async resetKey(key: string): Promise<void> {
159150
let resetQuery = `
160-
DELETE FROM rate_limit.individual_records
161-
WHERE key = $1 AND session_id = $2
151+
SELECT * FROM rate_limit.ind_reset_key($1, $2);
162152
`
163153

164154
try {
@@ -178,7 +168,7 @@ class PostgresStoreIndividualIP implements Store {
178168
*/
179169
async resetAll(): Promise<void> {
180170
let resetAllQuery = `
181-
DELETE FROM rate_limit.individual_records WHERE session_id = $1
171+
SELECT * FROM rate_limit.ind_reset_session($1);
182172
`
183173

184174
try {

‎source/util/session_handler.ts

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ export async function getSession(
2828
pool: Pool,
2929
): Promise<Session> {
3030
let selectSessionQuery =
31-
'SELECT id, name_, type_, expires_at FROM rate_limit.sessions WHERE name_ = $1 and type_ = $2 LIMIT 1'
31+
'SELECT id, name_, type_, expires_at FROM rate_limit.session_select($1, $2);'
3232
try {
3333
let result = await pool.query(selectSessionQuery, [name_, type_])
3434

@@ -87,24 +87,20 @@ export async function createNewSession(
8787
expires_at: calculateNextResetTime(windowMs),
8888
}
8989

90-
let deleteSessionQuery = `DELETE FROM rate_limit.sessions
91-
WHERE name_ = $1 and type_ = $2
92-
`
93-
let insertSessionQuery = `INSERT INTO rate_limit.sessions(name_, type_, expires_at)
94-
SELECT $1, $2, $3
95-
RETURNING id, name_, type_`
90+
let resetSessionQuery = `
91+
SELECT id, name_, type_ from rate_limit.session_reset($1, $2, $3)
92+
`
9693

97-
let client = await pool.connect()
9894
try {
99-
await client.query(deleteSessionQuery, [name_, type_])
100-
let result = await client.query(insertSessionQuery, [
95+
let result = await pool.query(resetSessionQuery, [
10196
name_,
10297
type_,
10398
newSession.expires_at,
10499
])
105100
newSession.id = result.rows[0].id
106-
} finally {
107-
client.release()
101+
} catch (err) {
102+
console.error(err)
103+
throw err
108104
}
109105

110106
return newSession

‎test/stores/store_aggregated_ip.spec.ts

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -71,19 +71,15 @@ describe('Postgres Store Aggregated IP', () => {
7171
let testStore = new PostgresStore({}, 'test')
7272
testStore.pool = pool
7373
testStore.session = newCreatedSession
74-
74+
let recordInsertGetRecordsQuery = `
75+
SELECT agg_increment as count FROM rate_limit.agg_increment($1, $2);
76+
`
7577
let incrementCount = await testStore.increment('key')
7678
sinon.assert.callCount(isSessionValidSpy, 1)
77-
sinon.assert.calledWith(
78-
query,
79-
`
80-
INSERT INTO rate_limit.records_aggregated(key, session_id) VALUES ($1, $2)
81-
ON CONFLICT ON CONSTRAINT unique_session_key DO UPDATE
82-
SET count = records_aggregated.count + 1
83-
RETURNING count
84-
`,
85-
['key', newCreatedSession.id],
86-
)
79+
sinon.assert.calledWith(query, recordInsertGetRecordsQuery, [
80+
'key',
81+
newCreatedSession.id,
82+
])
8783
assert.equal(incrementCount.totalHits, dbCount)
8884
assert.isTrue(
8985
(incrementCount.resetTime?.getMilliseconds() ||
@@ -106,9 +102,7 @@ describe('Postgres Store Aggregated IP', () => {
106102

107103
await testStore.decrement('key')
108104
let decrementQuery = `
109-
UPDATE rate_limit.records_aggregated
110-
SET count = greatest(0, count-1)
111-
WHERE key = $1 and session_id = $2
105+
SELECT * FROM rate_limit.agg_decrement($1, $2);
112106
`
113107
sinon.assert.calledWith(query, decrementQuery, [
114108
'key',
@@ -128,8 +122,7 @@ describe('Postgres Store Aggregated IP', () => {
128122

129123
await testStore.resetKey('key')
130124
let resetQuery = `
131-
DELETE FROM rate_limit.records_aggregated
132-
WHERE key = $1 and session_id = $2
125+
SELECT * FROM rate_limit.agg_reset_key($1, $2)
133126
`
134127
sinon.assert.calledWith(query, resetQuery, ['key', newCreatedSession.id])
135128
})
@@ -146,7 +139,7 @@ describe('Postgres Store Aggregated IP', () => {
146139

147140
await testStore.resetAll()
148141
let resetAllQuery = `
149-
DELETE FROM rate_limit.records_aggregated WHERE session_id = $1
142+
SELECT * FROM rate_limit.agg_reset_session($1);
150143
`
151144
sinon.assert.calledWith(query, resetAllQuery, [newCreatedSession.id])
152145
})

‎test/stores/store_individual_ip.spec.ts

Lines changed: 10 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -61,11 +61,8 @@ describe('Postgres Store Individual IP', () => {
6161
let dbCount = 1
6262

6363
isSessionValidSpy.returns(true)
64-
query.onFirstCall().returns({
65-
rows: [],
66-
})
6764

68-
query.onSecondCall().returns({
65+
query.onFirstCall().returns({
6966
rows: [
7067
{
7168
count: dbCount,
@@ -78,16 +75,12 @@ describe('Postgres Store Individual IP', () => {
7875

7976
let incrementCount = await testStore.increment('key')
8077
sinon.assert.callCount(isSessionValidSpy, 1)
81-
sinon.assert.calledWith(
82-
query,
83-
'INSERT INTO rate_limit.individual_records(key, session_id) VALUES ($1, $2)',
84-
['key', newCreatedSession.id],
85-
)
86-
sinon.assert.calledWith(
87-
query,
88-
'SELECT count(id) AS count FROM rate_limit.individual_records WHERE key = $1 AND session_id = $2',
89-
['key', newCreatedSession.id],
90-
)
78+
let recordInsertQuery =
79+
'SELECT ind_increment as count FROM rate_limit.ind_increment($1, $2)'
80+
sinon.assert.calledWith(query, recordInsertQuery, [
81+
'key',
82+
newCreatedSession.id,
83+
])
9184
assert.equal(incrementCount.totalHits, dbCount)
9285
assert.isTrue(
9386
(incrementCount.resetTime?.getMilliseconds() ||
@@ -110,13 +103,7 @@ describe('Postgres Store Individual IP', () => {
110103

111104
await testStore.decrement('key')
112105
let decrementQuery = `
113-
WITH
114-
rows_to_delete AS (
115-
SELECT id FROM rate_limit.individual_records
116-
WHERE key = $1 and session_id = $2 ORDER BY event_time LIMIT 1
117-
)
118-
DELETE FROM rate_limit.individual_records
119-
USING rows_to_delete WHERE individual_records.id = rows_to_delete.id
106+
SELECT * FROM rate_limit.ind_decrement($1, $2);
120107
`
121108
sinon.assert.calledWith(query, decrementQuery, [
122109
'key',
@@ -136,8 +123,7 @@ describe('Postgres Store Individual IP', () => {
136123

137124
await testStore.resetKey('key')
138125
let resetQuery = `
139-
DELETE FROM rate_limit.individual_records
140-
WHERE key = $1 AND session_id = $2
126+
SELECT * FROM rate_limit.ind_reset_key($1, $2);
141127
`
142128
sinon.assert.calledWith(query, resetQuery, ['key', newCreatedSession.id])
143129
})
@@ -154,7 +140,7 @@ describe('Postgres Store Individual IP', () => {
154140

155141
await testStore.resetAll()
156142
let resetAllQuery = `
157-
DELETE FROM rate_limit.individual_records WHERE session_id = $1
143+
SELECT * FROM rate_limit.ind_reset_session($1);
158144
`
159145
sinon.assert.calledWith(query, resetAllQuery, [newCreatedSession.id])
160146
})

0 commit comments

Comments
 (0)