From 378d6a5433748d5998e746682c5d2fa606817e0a Mon Sep 17 00:00:00 2001
From: pablo <p.lopez.lpz@gmail.com>
Date: Wed, 16 Jul 2025 09:22:29 +0300
Subject: [PATCH] fix: use kv store

---
 migrations/001_create_ratelimit_state.sql | 11 ++--
 ratelimit/store/sqlite.nim                | 74 ++++++++++++++---------
 tests/test_sqlite_store.nim               |  2 -
 3 files changed, 48 insertions(+), 39 deletions(-)

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