sqlite: adding connection pooling and WAL

.gitignore

@@ -64,4 +64,4 @@ cython_debug/
 config.yaml
 
 # SQLite
-*.db
+*.db*

db.py

@@ -16,9 +16,15 @@ class TokenDatabase:
         if not db_dir.exists():
             db_dir.mkdir(parents=True, exist_ok=True)
     
+    async def _get_connection(self):
+        """Return a connection with WAL mode enabled."""
+        conn= await aiosqlite.connect(self.db_path)
+        await conn.execute("PRAGMA journal_mode=WAL;")
+        return conn
+
     async def init_db(self):
         """Initialize the database tables."""
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         await db.execute('''
             CREATE TABLE IF NOT EXISTS token_counts (
                 endpoint TEXT,
@@ -46,7 +52,7 @@ class TokenDatabase:
     async def update_token_counts(self, endpoint: str, model: str, input_tokens: int, output_tokens: int):
         """Update token counts for a specific endpoint and model."""
         total_tokens = input_tokens + output_tokens
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         await db.execute('''
             INSERT INTO token_counts (endpoint, model, input_tokens, output_tokens, total_tokens)
             VALUES (?, ?, ?, ?, ?)
@@ -64,7 +70,7 @@ class TokenDatabase:
         now = datetime.now(tz=timezone.utc)
         timestamp = int(datetime(now.year, now.month, now.day, now.hour, now.minute).timestamp())
 
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         await db.execute('''
             INSERT INTO token_time_series (endpoint, model, input_tokens, output_tokens, total_tokens, timestamp)
             VALUES (?, ?, ?, ?, ?, ?)
@@ -75,7 +81,7 @@ class TokenDatabase:
         """Update multiple token counts in a single transaction."""
         if not counts:
             return
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         for endpoint, models in counts.items():
             for model, (input_tokens, output_tokens) in models.items():
                 total_tokens = input_tokens + output_tokens
@@ -92,7 +98,7 @@ class TokenDatabase:
 
     async def add_batched_time_series(self, entries: list):
         """Add multiple time series entries in a single transaction."""
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         for entry in entries:
             await db.execute('''
                 INSERT INTO token_time_series (endpoint, model, input_tokens, output_tokens, total_tokens, timestamp)
@@ -103,7 +109,7 @@ class TokenDatabase:
 
     async def load_token_counts(self):
         """Load all token counts from database."""
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         async with db.execute('SELECT endpoint, model, input_tokens, output_tokens, total_tokens FROM token_counts') as cursor:
             async for row in cursor:
                 yield {
@@ -116,7 +122,7 @@ class TokenDatabase:
 
     async def get_latest_time_series(self, limit: int = 100):
         """Get the latest time series entries."""
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         async with db.execute('''
             SELECT endpoint, model, input_tokens, output_tokens, total_tokens, timestamp
             FROM token_time_series
@@ -135,7 +141,7 @@ class TokenDatabase:
 
     async def get_token_counts_for_model(self, model):
         """Get token counts for a specific model, aggregated across all endpoints."""
-        async with aiosqlite.connect(self.db_path) as db:
+        db = await self._get_connection()
         async with db.execute('SELECT endpoint, model, input_tokens, output_tokens, total_tokens FROM token_counts WHERE model = ?', (model,)) as cursor:
             total_input = 0
             total_output = 0

static/index.html

@@ -122,7 +122,8 @@
             }
             .copy-link,
             .delete-link,
-            .show-link {
+            .show-link,
+            .stats-link {
                 font-size: 0.9em;
                 margin-left: 0.5em;
                 cursor: pointer;