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 )