fix: use kv store

This commit is contained in:
pablo 2025-07-16 09:22:29 +03:00
parent 7e4f930ae3
commit 378d6a5433
No known key found for this signature in database
GPG Key ID: 78F35FCC60FDC63A
3 changed files with 48 additions and 39 deletions

View File

@ -1,7 +1,4 @@
-- will only exist one row in the table
CREATE TABLE IF NOT EXISTS bucket_state (
id INTEGER PRIMARY KEY,
budget INTEGER NOT NULL,
budget_cap INTEGER NOT NULL,
last_time_full_seconds INTEGER NOT NULL
)
CREATE TABLE IF NOT EXISTS kv_store (
key TEXT PRIMARY KEY,
value BLOB
);

View File

@ -1,5 +1,6 @@
import std/times
import std/strutils
import std/json
import ./store
import chronos
import db_connector/db_sqlite
@ -9,6 +10,8 @@ type SqliteRateLimitStore* = ref object
db: DbConn
dbPath: string
const BUCKET_STATE_KEY = "rate_limit_bucket_state"
proc newSqliteRateLimitStore*(dbPath: string = ":memory:"): SqliteRateLimitStore =
result = SqliteRateLimitStore(dbPath: dbPath)
result.db = open(dbPath, "", "", "")
@ -16,25 +19,35 @@ proc newSqliteRateLimitStore*(dbPath: string = ":memory:"): SqliteRateLimitStore
# Create table if it doesn't exist
result.db.exec(
sql"""
CREATE TABLE IF NOT EXISTS bucket_state (
id INTEGER PRIMARY KEY,
budget INTEGER NOT NULL,
budget_cap INTEGER NOT NULL,
last_time_full_seconds INTEGER NOT NULL
CREATE TABLE IF NOT EXISTS kv_store (
key TEXT PRIMARY KEY,
value BLOB
)
"""
)
# Insert default state if table is empty
let count = result.db.getValue(sql"SELECT COUNT(*) FROM bucket_state").parseInt()
# Insert default state if key doesn't exist
let count = result.db
.getValue(sql"SELECT COUNT(*) FROM kv_store WHERE key = ?", BUCKET_STATE_KEY)
.parseInt()
if count == 0:
let defaultTimeSeconds = Moment.now().epochSeconds()
let defaultState = BucketState(
budget: 10,
budgetCap: 10,
lastTimeFull: Moment.init(defaultTimeSeconds, chronos.seconds(1)),
)
# Serialize to JSON
let jsonState =
%*{
"budget": defaultState.budget,
"budgetCap": defaultState.budgetCap,
"lastTimeFullSeconds": defaultTimeSeconds,
}
result.db.exec(
sql"""
INSERT INTO bucket_state (id, budget, budget_cap, last_time_full_seconds)
VALUES (1, 10, 10, ?)
""",
defaultTimeSeconds,
sql"INSERT INTO kv_store (key, value) VALUES (?, ?)", BUCKET_STATE_KEY, $jsonState
)
proc close*(store: SqliteRateLimitStore) =
@ -47,32 +60,33 @@ proc saveBucketState*(
try:
# Convert Moment to Unix seconds for storage
let lastTimeSeconds = bucketState.lastTimeFull.epochSeconds()
# Serialize to JSON
let jsonState =
%*{
"budget": bucketState.budget,
"budgetCap": bucketState.budgetCap,
"lastTimeFullSeconds": lastTimeSeconds,
}
store.db.exec(
sql"""
UPDATE bucket_state
SET budget = ?, budget_cap = ?, last_time_full_seconds = ?
WHERE id = 1
""",
bucketState.budget,
bucketState.budgetCap,
lastTimeSeconds,
sql"UPDATE kv_store SET value = ? WHERE key = ?", $jsonState, BUCKET_STATE_KEY
)
return true
except:
return false
proc loadBucketState*(store: SqliteRateLimitStore): Future[BucketState] {.async.} =
let row = store.db.getRow(
sql"""
SELECT budget, budget_cap, last_time_full_seconds
FROM bucket_state
WHERE id = 1
"""
)
# Convert Unix seconds back to Moment (seconds precission)
let unixSeconds = row[2].parseInt().int64
let jsonStr =
store.db.getValue(sql"SELECT value FROM kv_store WHERE key = ?", BUCKET_STATE_KEY)
# Parse JSON and reconstruct BucketState
let jsonData = parseJson(jsonStr)
let unixSeconds = jsonData["lastTimeFullSeconds"].getInt().int64
let lastTimeFull = Moment.init(unixSeconds, chronos.seconds(1))
return BucketState(
budget: row[0].parseInt(), budgetCap: row[1].parseInt(), lastTimeFull: lastTimeFull
budget: jsonData["budget"].getInt(),
budgetCap: jsonData["budgetCap"].getInt(),
lastTimeFull: lastTimeFull,
)

View File

@ -1,5 +1,3 @@
{.used.}
import testutils/unittests
import ../ratelimit/store/sqlite
import ../ratelimit/store/store