Skip to content

Commit a36c775

Browse files
feat: moved session to database stored procedure for individual store
1 parent 723dc32 commit a36c775

File tree

2 files changed

+119
-27
lines changed

2 files changed

+119
-27
lines changed
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
--Migration generated Sat, 17 Feb 2024 15:02:12 +0000
2+
3+
DROP FUNCTION IF EXISTS rate_limit.ind_increment(key_ text, session_id_ uuid);
4+
DROP FUNCTION IF EXISTS rate_limit.ind_decrement(key_ text, session_id_ uuid);
5+
DROP FUNCTION IF EXISTS rate_limit.ind_reset_key(key_ text, session_id_ uuid);
6+
DROP FUNCTION IF EXISTS rate_limit.ind_reset_session(session_id_ uuid);
7+
8+
CREATE OR REPLACE FUNCTION rate_limit.ind_increment(key_ text, prefix text, window_ms double precision, reference_time timestamptz DEFAULT now())
9+
RETURNS record AS
10+
$bd$
11+
DECLARE
12+
in_session_id uuid;
13+
in_session_expiration timestamptz;
14+
session_type text = 'individual';
15+
record_count int = 0;
16+
ret RECORD;
17+
BEGIN
18+
19+
LOCK TABLE rate_limit.sessions;
20+
21+
SELECT id, expires_at
22+
FROM rate_limit.session_select($2, session_type)
23+
WHERE expires_at > $4
24+
INTO in_session_id, in_session_expiration;
25+
26+
IF in_session_id is null THEN
27+
in_session_expiration = to_timestamp(extract (epoch from $4)+ $3/1000.0);
28+
SELECT id, in_session_expiration
29+
FROM rate_limit.session_reset(
30+
$2, session_type, in_session_expiration
31+
)
32+
INTO in_session_id;
33+
END IF;
34+
35+
36+
INSERT INTO rate_limit.individual_records(key, session_id) VALUES ($1, in_session_id);
37+
38+
SELECT count(id)::int AS count FROM rate_limit.individual_records WHERE key = $1 AND session_id = in_session_id
39+
INTO record_count;
40+
41+
ret:= (record_count, in_session_expiration);
42+
43+
RETURN ret;
44+
END;
45+
$bd$
46+
LANGUAGE plpgsql;
47+
48+
CREATE OR REPLACE FUNCTION rate_limit.ind_decrement(key_ text, prefix text, reference_time timestamptz DEFAULT now())
49+
RETURNS void AS
50+
$bd$
51+
DECLARE
52+
in_session_id uuid;
53+
session_type text = 'individual';
54+
BEGIN
55+
56+
SELECT id
57+
FROM rate_limit.session_select($2, session_type)
58+
WHERE expires_at > $3
59+
INTO in_session_id;
60+
61+
WITH
62+
rows_to_delete AS (
63+
SELECT id FROM rate_limit.individual_records
64+
WHERE key = $1 and session_id = in_session_id ORDER BY event_time LIMIT 1
65+
)
66+
DELETE FROM rate_limit.individual_records
67+
USING rows_to_delete WHERE individual_records.id = rows_to_delete.id;
68+
END;
69+
$bd$
70+
LANGUAGE plpgsql;
71+
72+
CREATE OR REPLACE FUNCTION rate_limit.ind_reset_key(key_ text, prefix text, reference_time timestamptz DEFAULT now())
73+
RETURNS void AS
74+
$bd$
75+
DECLARE
76+
in_session_id uuid;
77+
session_type text = 'individual';
78+
BEGIN
79+
80+
SELECT id
81+
FROM rate_limit.session_select($2, session_type)
82+
WHERE expires_at > $3
83+
INTO in_session_id;
84+
85+
DELETE FROM rate_limit.individual_records
86+
WHERE key = $1 AND session_id = in_session_id;
87+
END;
88+
$bd$
89+
LANGUAGE plpgsql;
90+
91+
CREATE OR REPLACE FUNCTION rate_limit.ind_reset_session(prefix text, reference_time timestamptz DEFAULT now())
92+
RETURNS void AS
93+
$bd$
94+
DECLARE
95+
in_session_id uuid;
96+
session_type text = 'individual';
97+
BEGIN
98+
99+
SELECT id
100+
FROM rate_limit.session_select($1, session_type)
101+
WHERE expires_at > $2
102+
INTO in_session_id;
103+
104+
DELETE FROM rate_limit.individual_records
105+
WHERE session_id = in_session_id;
106+
END;
107+
$bd$
108+
LANGUAGE plpgsql;

source/stores/individual_ip/store_individual_ip.ts

Lines changed: 11 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
import { Store, Options, ClientRateLimitInfo } from 'express-rate-limit'
22

33
import Pool from 'pg-pool'
4-
import { Session } from '../../models/session'
5-
import { getSession, isSessionValid } from '../../util/session_handler'
64
import { applyMigrations } from '../../util/migration_handler'
75

86
/**
@@ -34,16 +32,6 @@ class PostgresStoreIndividualIP implements Store {
3432
*/
3533
SESSION_TYPE: string = 'individual'
3634

37-
/**
38-
* The session instance persisted on the client side.
39-
*/
40-
session: Session = {
41-
id: 'init',
42-
name: 'init',
43-
type: this.SESSION_TYPE,
44-
expires_at: undefined,
45-
}
46-
4735
/**
4836
* The duration of time before which all hit counts are reset (in milliseconds).
4937
*/
@@ -91,24 +79,20 @@ class PostgresStoreIndividualIP implements Store {
9179
*/
9280
async increment(key: string): Promise<ClientRateLimitInfo> {
9381
let recordInsertQuery =
94-
'SELECT ind_increment as count FROM rate_limit.ind_increment($1, $2)'
95-
if (!isSessionValid(this.session)) {
96-
this.session = await getSession(
97-
this.prefix,
98-
this.SESSION_TYPE,
99-
this.windowMs,
100-
this.pool,
101-
)
102-
}
82+
'SELECT * FROM rate_limit.ind_increment($1, $2, $3) AS (count int, expires_at timestamptz);'
10383

10484
try {
10585
let result = await this.pool.query(recordInsertQuery, [
10686
key,
107-
this.session.id,
87+
this.prefix,
88+
this.windowMs,
10889
])
10990
let totalHits: number = 0
110-
if (result.rows.length > 0) totalHits = parseInt(result.rows[0].count)
111-
let resetTime: Date | undefined = this.session.expires_at
91+
let resetTime: Date | undefined = undefined
92+
if (result.rows.length > 0) {
93+
totalHits = parseInt(result.rows[0].count)
94+
resetTime = result.rows[0].expires_at
95+
}
11296
return {
11397
totalHits,
11498
resetTime,
@@ -132,7 +116,7 @@ class PostgresStoreIndividualIP implements Store {
132116
`
133117

134118
try {
135-
await this.pool.query(decrementQuery, [key, this.session.id])
119+
await this.pool.query(decrementQuery, [key, this.prefix])
136120
} catch (err) {
137121
console.error(err)
138122
throw err
@@ -152,7 +136,7 @@ class PostgresStoreIndividualIP implements Store {
152136
`
153137

154138
try {
155-
await this.pool.query(resetQuery, [key, this.session.id])
139+
await this.pool.query(resetQuery, [key, this.prefix])
156140
} catch (err) {
157141
console.error(err)
158142
throw err
@@ -172,7 +156,7 @@ class PostgresStoreIndividualIP implements Store {
172156
`
173157

174158
try {
175-
await this.pool.query(resetAllQuery, [this.session.id])
159+
await this.pool.query(resetAllQuery, [this.prefix])
176160
} catch (err) {
177161
console.error(err)
178162
throw err

0 commit comments

Comments
 (0)