luckyimnotanigger/LibreTranslate
1
0
Fork 0
mirror of https://github.com/LibreTranslate/LibreTranslate.git synced 2025-06-19 07:30:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Granular API key char limit support

Browse source
This commit is contained in:
Piero Toffanin 2024-01-21 12:48:42 -05:00
parent 1d8b58e2bf
commit 2b11e909d3
5 changed files with 66 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

46
libretranslate/api_keys.py
View file

@ -32,41 +32,49 @@ class Database:
"""CREATE TABLE IF NOT EXISTS api_keys (
"api_key" TEXT NOT NULL,
"req_limit" INTEGER NOT NULL,
"char_limit" INTEGER DEFAULT NULL,
PRIMARY KEY("api_key")
);"""
)
# Schema/upgrade checks
schema = self.c.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='api_keys';").fetchone()[0]
if '"char_limit" INTEGER DEFAULT NULL' not in schema:
self.c.execute('ALTER TABLE api_keys ADD COLUMN "char_limit" INTEGER DEFAULT NULL;')
def lookup(self, api_key):
req_limit = self.cache.get(api_key)
if req_limit is None:
val = self.cache.get(api_key)
if val is None:
# DB Lookup
stmt = self.c.execute(
"SELECT req_limit FROM api_keys WHERE api_key = ?", (api_key,)
"SELECT req_limit, char_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]
self.cache[api_key] = row
val = row
else:
self.cache[api_key] = False
req_limit = False
val = False
if isinstance(req_limit, bool):
req_limit = None
if isinstance(val, bool):
val = None
return req_limit
return val
def add(self, req_limit, api_key="auto"):
def add(self, req_limit, api_key="auto", char_limit=None):
if api_key == "auto":
api_key = str(uuid.uuid4())
if char_limit == 0:
char_limit = None
self.remove(api_key)
self.c.execute(
"INSERT INTO api_keys (api_key, req_limit) VALUES (?, ?)",
(api_key, req_limit),
"INSERT INTO api_keys (api_key, req_limit, char_limit) VALUES (?, ?, ?)",
(api_key, req_limit, char_limit),
)
self.c.commit()
return (api_key, req_limit)
return (api_key, req_limit, char_limit)
def remove(self, api_key):
self.c.execute("DELETE FROM api_keys WHERE api_key = ?", (api_key,))
@ -74,7 +82,7 @@ class Database:
return api_key
def all(self):
row = self.c.execute("SELECT api_key, req_limit FROM api_keys")
row = self.c.execute("SELECT api_key, req_limit, char_limit FROM api_keys")
return row.fetchall()
@ -84,8 +92,8 @@ class RemoteDatabase:
self.cache = ExpiringDict(max_len=max_cache_len, max_age_seconds=max_cache_age)
def lookup(self, api_key):
req_limit = self.cache.get(api_key)
if req_limit is None:
val = self.cache.get(api_key)
if val is None:
try:
r = requests.post(self.url, data={'api_key': api_key}, timeout=60)
res = r.json()
@ -94,6 +102,8 @@ class RemoteDatabase:
return None
req_limit = res.get('req_limit', None) if res.get('error', None) is None else None
self.cache[api_key] = req_limit
char_limit = res.get('char_limit', None) if res.get('error', None) is None else None
return req_limit
self.cache[api_key] = (req_limit, char_limit)
return val