110 lines
4.9 KiB
Python
110 lines
4.9 KiB
Python
#!/usr/bin/env python3.12
|
|
# pylint: disable=bare-except, broad-exception-raised
|
|
|
|
"""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):
|
|
"""Global State for API"""
|
|
def __init__(self, app: FastAPI, util, constants): # pylint: disable=unused-argument
|
|
super().__init__()
|
|
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"
|
|
await _db.executescript("pragma journal_mode = WAL; pragma synchronous = normal; pragma temp_store = memory; pragma mmap_size = 30000000000;")
|
|
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 not self.counters_initialized:
|
|
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)
|
|
|
|
await _db.executescript("pragma journal_mode = WAL; pragma synchronous = normal; pragma temp_store = memory; pragma mmap_size = 30000000000;")
|
|
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):
|
|
"""Increment Counter"""
|
|
if not counter in self.counters.keys():
|
|
raise BaseException(f"[State] Counter {counter} does not exist")
|
|
|
|
self.counters[counter] += 1
|
|
return True
|
|
|
|
async def get_counter(self, counter: str):
|
|
"""Get Counter"""
|
|
if not counter in self.counters.keys():
|
|
raise BaseException(f"[State] Counter {counter} does not exist")
|
|
|
|
return self.counters[counter]
|
|
|
|
async def get_all_counters(self):
|
|
"""Get All Counters"""
|
|
return self.counters |