api/state.py

#!/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