#!/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("/", "usr", "local", "share", "sqlite_dbs", "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