File size: 2,567 Bytes
0cd599e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import sqlite3
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator
def utc_now() -> str:
return datetime.now(timezone.utc).isoformat(timespec="seconds")
class Database:
def __init__(self, path: Path):
self.path = path
def open_connection(self) -> sqlite3.Connection:
connection = sqlite3.connect(self.path, timeout=10)
try:
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA foreign_keys=ON")
connection.execute("PRAGMA busy_timeout=10000")
except sqlite3.Error:
connection.close()
raise
return connection
@contextmanager
def connect(self) -> Iterator[sqlite3.Connection]:
connection = self.open_connection()
try:
yield connection
connection.commit()
except Exception:
connection.rollback()
raise
finally:
connection.close()
def initialize(self) -> None:
self.path.parent.mkdir(parents=True, exist_ok=True)
with self.connect() as db:
db.execute("PRAGMA journal_mode=WAL")
db.executescript("""
CREATE TABLE IF NOT EXISTS sessions (
digest TEXT PRIMARY KEY, csrf_digest TEXT NOT NULL,
expires_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS login_attempts (
source TEXT NOT NULL, username TEXT NOT NULL,
attempted_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS login_attempts_time ON login_attempts(attempted_at);
CREATE TABLE IF NOT EXISTS api_keys (
id INTEGER PRIMARY KEY, name TEXT NOT NULL, digest TEXT NOT NULL UNIQUE,
mask TEXT NOT NULL, created_at TEXT NOT NULL, revoked_at TEXT,
last_used_at TEXT
);
CREATE TABLE IF NOT EXISTS usage_events (
id INTEGER PRIMARY KEY, created_at TEXT NOT NULL,
source TEXT NOT NULL, key_id INTEGER REFERENCES api_keys(id),
model TEXT NOT NULL, input_tokens INTEGER NOT NULL,
output_tokens INTEGER NOT NULL, duration_ms INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS usage_events_date ON usage_events(created_at);
""")
|