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

import std/times
import std/strutils
import ./store
import chronos
import db_connector/db_sqlite

# SQLite Implementation
type SqliteRateLimitStore* = ref object
  db: DbConn
  dbPath: string

proc newSqliteRateLimitStore*(dbPath: string = ":memory:"): SqliteRateLimitStore =
  result = SqliteRateLimitStore(dbPath: dbPath)
  result.db = open(dbPath, "", "", "")

  # 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
    )
  """
  )

  # Insert default state if table is empty
  let count = result.db.getValue(sql"SELECT COUNT(*) FROM bucket_state").parseInt()
  if count == 0:
    let defaultTimeSeconds = Moment.now().epochSeconds()
    result.db.exec(
      sql"""
      INSERT INTO bucket_state (id, budget, budget_cap, last_time_full_seconds) 
      VALUES (1, 10, 10, ?)
    """,
      defaultTimeSeconds,
    )

proc close*(store: SqliteRateLimitStore) =
  if store.db != nil:
    store.db.close()

proc saveBucketState*(
    store: SqliteRateLimitStore, bucketState: BucketState
): Future[bool] {.async.} =
  try:
    # Convert Moment to Unix seconds for storage
    let lastTimeSeconds = bucketState.lastTimeFull.epochSeconds()
    store.db.exec(
      sql"""
      UPDATE bucket_state 
      SET budget = ?, budget_cap = ?, last_time_full_seconds = ?
      WHERE id = 1
    """,
      bucketState.budget,
      bucketState.budgetCap,
      lastTimeSeconds,
    )
    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 lastTimeFull = Moment.init(unixSeconds, chronos.seconds(1))

  return BucketState(
    budget: row[0].parseInt(), budgetCap: row[1].parseInt(), lastTimeFull: lastTimeFull
  )