import sqlite3 import uuid from expiringdict import ExpiringDict DEFAULT_DB_PATH = "api_keys.db" class Database: def __init__(self, db_path=DEFAULT_DB_PATH, max_cache_len=1000, max_cache_age=30): self.db_path = db_path self.cache = ExpiringDict(max_len=max_cache_len, max_age_seconds=max_cache_age) # Make sure to do data synchronization on writes! self.c = sqlite3.connect(db_path, check_same_thread=False) self.c.execute( """CREATE TABLE IF NOT EXISTS api_keys ( "api_key" TEXT NOT NULL, "req_limit" INTEGER NOT NULL, PRIMARY KEY("api_key") );""" ) def lookup(self, api_key): req_limit = self.cache.get(api_key) if req_limit is None: # DB Lookup stmt = self.c.execute( "SELECT req_limit FROM api_keys WHERE api_key = ?", (api_key,) ) row = stmt.fetchone() if row is not None: self.cache[api_key] = row[0] req_limit = row[0] else: self.cache[api_key] = False req_limit = False if isinstance(req_limit, bool): req_limit = None return req_limit def add(self, req_limit, api_key="auto"): if api_key == "auto": api_key = str(uuid.uuid4()) self.remove(api_key) self.c.execute( "INSERT INTO api_keys (api_key, req_limit) VALUES (?, ?)", (api_key, req_limit), ) self.c.commit() return (api_key, req_limit) def remove(self, api_key): self.c.execute("DELETE FROM api_keys WHERE api_key = ?", (api_key,)) self.c.commit() return api_key def all(self): row = self.c.execute("SELECT api_key, req_limit FROM api_keys") return row.fetchall()