nim-chat-sdk/ratelimit/store/sqlite.nim

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,
    )
  )