110 lines
4.4 KiB
Python
110 lines
4.4 KiB
Python
|
#!/usr/bin/env python3.12
|
||
|
|
|
||
|
|
"""Global State Storage/Counters"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import aiosqlite as sqlite3
|
||
|
|
|
||
|
|
from fastapi import FastAPI
|
||
|
|
from fastapi_utils.tasks import repeat_every
|
||
|
|
|
||
|
|
|
||
|
|
class State(FastAPI):
|
||
|
|
def __init__(self, app: FastAPI, util, constants):
|
||
|
|
self.counter_db_path = os.path.join("/", "var", "lib", "singerdbs", "stats.db")
|
||
|
|
self.counters = {
|
||
|
|
str(counter): 0 for counter in constants.AVAILABLE_COUNTERS
|
||
|
|
}
|
||
|
|
self.counters_initialized = False
|
||
|
|
logging.debug("[State] Counters: %s", self.counters)
|
||
|
|
|
||
|
|
@app.on_event("startup")
|
||
|
|
async def get_counters():
|
||
|
|
logging.info("[State] Initializing counters...")
|
||
|
|
async with sqlite3.connect(self.counter_db_path, timeout=2) as _db:
|
||
|
|
_query = "SELECT ai_requests, lyric_requests, transcript_list_requests, transcript_requests, lyrichistory_requests, \
|
||
|
|
failedlyric_requests, misc_failures, claude_ai_requests FROM counters LIMIT 1"
|
||
|
|
async with _db.execute(_query) as _cursor:
|
||
|
|
_result = await _cursor.fetchone()
|
||
|
|
(ai_requests,
|
||
|
|
lyric_requests,
|
||
|
|
transcript_list_requests,
|
||
|
|
transcript_requests,
|
||
|
|
lyrichistory_requests,
|
||
|
|
failedlyric_requests,
|
||
|
|
misc_failures,
|
||
|
|
claude_ai_requests) = _result
|
||
|
|
self.counters = {
|
||
|
|
'ai_requests': ai_requests,
|
||
|
|
'lyric_requests': lyric_requests,
|
||
|
|
'transcript_list_requests': transcript_list_requests,
|
||
|
|
'transcript_requests': transcript_requests,
|
||
|
|
'lyrichistory_requests': lyrichistory_requests,
|
||
|
|
'failedlyric_requests': failedlyric_requests,
|
||
|
|
'misc_failures': misc_failures,
|
||
|
|
'claude_ai_requests': claude_ai_requests
|
||
|
|
}
|
||
|
|
self.counters_initialized = True
|
||
|
|
logging.info("Counters loaded from db: %s", self.counters)
|
||
|
|
|
||
|
|
|
||
|
|
@app.on_event("startup")
|
||
|
|
@repeat_every(seconds=10)
|
||
|
|
async def update_db():
|
||
|
|
if self.counters_initialized == False:
|
||
|
|
logging.debug("[State] TICK: Counters not yet initialized")
|
||
|
|
return
|
||
|
|
|
||
|
|
ai_requests = self.counters.get('ai_requests')
|
||
|
|
lyric_requests = self.counters.get('lyric_requests')
|
||
|
|
transcript_list_requests = self.counters.get('transcript_list_requests')
|
||
|
|
transcript_requests = self.counters.get('transcript_requests')
|
||
|
|
lyrichistory_requests = self.counters.get('lyrichistory_requests')
|
||
|
|
failedlyric_requests = self.counters.get('failedlyric_requests')
|
||
|
|
claude_ai_requests = self.counters.get('claude_ai_requests')
|
||
|
|
|
||
|
|
async with sqlite3.connect(self.counter_db_path, timeout=2) as _db:
|
||
|
|
_query = "UPDATE counters SET ai_requests = ?, lyric_requests = ?, transcript_list_requests = ?, \
|
||
|
|
transcript_requests = ?, lyrichistory_requests = ?, failedlyric_requests = ?, \
|
||
|
|
claude_ai_requests = ?"
|
||
|
|
|
||
|
|
_params = (ai_requests,
|
||
|
|
lyric_requests,
|
||
|
|
transcript_list_requests,
|
||
|
|
transcript_requests,
|
||
|
|
lyrichistory_requests,
|
||
|
|
failedlyric_requests,
|
||
|
|
claude_ai_requests)
|
||
|
|
async with _db.execute(_query, _params) as _cursor:
|
||
|
|
if _cursor.rowcount != 1:
|
||
|
|
logging.error("Failed to update DB")
|
||
|
|
return
|
||
|
|
await _db.commit()
|
||
|
|
logging.debug("[State] Updated DB")
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
async def increment_counter(self, counter: str):
|
||
|
|
if not(counter in self.counters.keys()):
|
||
|
|
raise BaseException("[State] Counter %s does not exist", counter)
|
||
|
|
|
||
|
|
self.counters[counter] += 1
|
||
|
|
return True
|
||
|
|
|
||
|
|
async def get_counter(self, counter: str):
|
||
|
|
if not(counter in self.counters.keys()):
|
||
|
|
raise BaseException("[State] Counter %s does not exist", counter)
|
||
|
|
|
||
|
|
return self.counters[counter]
|
||
|
|
|
||
|
|
async def get_all_counters(self):
|
||
|
|
return self.counters
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|