Skip to content

Commit 492836e

Browse files
adrianprelipceanadrianprelipcean
committed
feat: moved session to database stored procedure for aggregated store
1 parent cdb48dd commit 492836e

File tree

2 files changed

+117
-29
lines changed

2 files changed

+117
-29
lines changed

‎source/migrations/6-move-session-to-db-agg.sql

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
--Migration generated Sat, 17 Feb 2024 14:02:12 +0000
2+
3+
DROP FUNCTION IF EXISTS rate_limit.agg_increment(key_ text, session_id_ uuid);
4+
DROP FUNCTION IF EXISTS rate_limit.agg_decrement(key_ text, session_id_ uuid);
5+
DROP FUNCTION IF EXISTS rate_limit.agg_reset_key(key_ text, session_id_ uuid);
6+
DROP FUNCTION IF EXISTS rate_limit.agg_reset_session(session_id_ uuid);
7+
8+
CREATE OR REPLACE FUNCTION rate_limit.agg_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 = 'aggregated';
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.records_aggregated(key, session_id)
37+
VALUES ($1, in_session_id)
38+
ON CONFLICT ON CONSTRAINT unique_session_key DO UPDATE
39+
SET count = records_aggregated.count + 1
40+
RETURNING count INTO record_count;
41+
42+
ret:= (record_count, in_session_expiration);
43+
44+
RETURN ret;
45+
END;
46+
$bd$
47+
LANGUAGE plpgsql;
48+
49+
CREATE OR REPLACE FUNCTION rate_limit.agg_decrement(key_ text, prefix text, reference_time timestamptz DEFAULT now())
50+
RETURNS void AS
51+
$bd$
52+
DECLARE
53+
in_session_id uuid;
54+
session_type text = 'aggregated';
55+
BEGIN
56+
57+
select id
58+
FROM rate_limit.session_select($2, session_type)
59+
WHERE expires_at > $3
60+
INTO in_session_id;
61+
62+
UPDATE rate_limit.records_aggregated
63+
SET count = greatest(0, count-1)
64+
WHERE key = $1 and session_id = in_session_id;
65+
END;
66+
$bd$
67+
LANGUAGE plpgsql;
68+
69+
CREATE OR REPLACE FUNCTION rate_limit.agg_reset_key(key_ text, prefix text, reference_time timestamptz DEFAULT now())
70+
RETURNS void AS
71+
$bd$
72+
DECLARE
73+
in_session_id uuid;
74+
session_type text = 'aggregated';
75+
BEGIN
76+
77+
SELECT id
78+
FROM rate_limit.session_select($2, session_type)
79+
WHERE expires_at > $3
80+
INTO in_session_id;
81+
82+
DELETE FROM rate_limit.records_aggregated
83+
WHERE key = $1 and session_id = in_session_id;
84+
END;
85+
$bd$
86+
LANGUAGE plpgsql;
87+
88+
CREATE OR REPLACE FUNCTION rate_limit.agg_reset_session(prefix text, reference_time timestamptz DEFAULT now())
89+
RETURNS void AS
90+
$bd$
91+
DECLARE
92+
in_session_id uuid;
93+
session_type text = 'aggregated';
94+
BEGIN
95+
96+
SELECT id
97+
FROM rate_limit.session_select($1, session_type)
98+
WHERE expires_at > $2
99+
INTO in_session_id;
100+
101+
DELETE FROM rate_limit.records_aggregated
102+
WHERE session_id = in_session_id;
103+
END;
104+
$bd$
105+
LANGUAGE plpgsql;

‎source/stores/aggregated_ip/store_aggregated_ip.ts

Lines changed: 12 additions & 29 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 PostgresStoreAggregatedIP implements Store {
3432
*/
3533
pool: any
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
*/
@@ -90,26 +78,21 @@ class PostgresStoreAggregatedIP implements Store {
9078
* @public
9179
*/
9280
async increment(key: string): Promise<ClientRateLimitInfo> {
93-
let recordInsertGetRecordsQuery = `
94-
SELECT agg_increment as count FROM rate_limit.agg_increment($1, $2);
95-
`
96-
if (!isSessionValid(this.session)) {
97-
this.session = await getSession(
98-
this.prefix,
99-
this.SESSION_TYPE,
100-
this.windowMs,
101-
this.pool,
102-
)
103-
}
81+
let recordInsertGetRecordsQuery = `SELECT * FROM rate_limit.agg_increment($1, $2, $3) AS (count int, expires_at timestamptz);`
10482

10583
try {
10684
let result = await this.pool.query(recordInsertGetRecordsQuery, [
10785
key,
108-
this.session.id,
86+
this.prefix,
87+
this.windowMs,
10988
])
11089
let totalHits: number = 0
111-
if (result.rows.length > 0) totalHits = parseInt(result.rows[0].count)
112-
let resetTime: Date | undefined = this.session.expires_at
90+
let resetTime: Date | undefined = undefined
91+
92+
if (result.rows.length > 0) {
93+
totalHits = parseInt(result.rows[0].count)
94+
resetTime = result.rows[0].expires_at
95+
}
11396
return {
11497
totalHits,
11598
resetTime,
@@ -133,7 +116,7 @@ class PostgresStoreAggregatedIP implements Store {
133116
`
134117

135118
try {
136-
await this.pool.query(decrementQuery, [key, this.session.id])
119+
await this.pool.query(decrementQuery, [key, this.prefix])
137120
} catch (err) {
138121
console.error(err)
139122
throw err
@@ -153,7 +136,7 @@ class PostgresStoreAggregatedIP implements Store {
153136
`
154137

155138
try {
156-
await this.pool.query(resetQuery, [key, this.session.id])
139+
await this.pool.query(resetQuery, [key, this.prefix])
157140
} catch (err) {
158141
console.error(err)
159142
throw err
@@ -173,7 +156,7 @@ class PostgresStoreAggregatedIP implements Store {
173156
`
174157

175158
try {
176-
await this.pool.query(resetAllQuery, [this.session.id])
159+
await this.pool.query(resetAllQuery, [this.prefix])
177160
} catch (err) {
178161
console.error(err)
179162
throw err

0 commit comments

Comments
 (0)