diff --git a/migrations/001_create_ratelimit_state.sql b/migrations/001_create_ratelimit_state.sql index 11431ab..030377c 100644 --- a/migrations/001_create_ratelimit_state.sql +++ b/migrations/001_create_ratelimit_state.sql @@ -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 - ) \ No newline at end of file +CREATE TABLE IF NOT EXISTS kv_store ( + key TEXT PRIMARY KEY, + value BLOB +); \ No newline at end of file diff --git a/ratelimit/store/sqlite.nim b/ratelimit/store/sqlite.nim index d8aaf38..ba3e72e 100644 --- a/ratelimit/store/sqlite.nim +++ b/ratelimit/store/sqlite.nim @@ -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, ) diff --git a/tests/test_sqlite_store.nim b/tests/test_sqlite_store.nim index 98dbe08..8728992 100644 --- a/tests/test_sqlite_store.nim +++ b/tests/test_sqlite_store.nim @@ -1,5 +1,3 @@ -{.used.} - import testutils/unittests import ../ratelimit/store/sqlite import ../ratelimit/store/store