import std/[times, strutils, json, options] import ./store import chronos import db_connector/db_sqlite # SQLite Implementation type SqliteRateLimitStore* = ref object db: DbConn dbPath: string const BUCKET_STATE_KEY = "rate_limit_bucket_state" proc newSqliteRateLimitStore*(db: DbConn): SqliteRateLimitStore = result = SqliteRateLimitStore(db: db) proc saveBucketState*( store: SqliteRateLimitStore, bucketState: BucketState ): Future[bool] {.async.} = try: # Convert Moment to Unix seconds for storage let lastTimeSeconds = bucketState.lastTimeFull.epochSeconds() let jsonState = %*{ "budget": bucketState.budget, "budgetCap": bucketState.budgetCap, "lastTimeFullSeconds": lastTimeSeconds, } store.db.exec( sql"INSERT INTO kv_store (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", BUCKET_STATE_KEY, $jsonState, ) return true except: return false proc loadBucketState*( store: SqliteRateLimitStore ): Future[Option[BucketState]] {.async.} = let jsonStr = store.db.getValue(sql"SELECT value FROM kv_store WHERE key = ?", BUCKET_STATE_KEY) if jsonStr == "": return none(BucketState) let jsonData = parseJson(jsonStr) let unixSeconds = jsonData["lastTimeFullSeconds"].getInt().int64 let lastTimeFull = Moment.init(unixSeconds, chronos.seconds(1)) return some( BucketState( budget: jsonData["budget"].getInt(), budgetCap: jsonData["budgetCap"].getInt(), lastTimeFull: lastTimeFull, ) )