mirror of
https://github.com/logos-messaging/nim-chat-sdk.git
synced 2026-01-03 22:53:12 +00:00
fix: use kv store
This commit is contained in:
parent
7e4f930ae3
commit
378d6a5433
@ -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
|
||||
);
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -1,5 +1,3 @@
|
||||
{.used.}
|
||||
|
||||
import testutils/unittests
|
||||
import ../ratelimit/store/sqlite
|
||||
import ../ratelimit/store/store
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user