59 lines
1.6 KiB
Python

import sqlite3
from contextlib import contextmanager
from datetime import datetime
from typing import Generator
from massranker.state import StateT
class Storage(object):
def __init__(self, db: sqlite3.Connection):
self._db = db
@classmethod
@contextmanager
def transact(cls, name) -> Generator["Storage", None, None]:
db = sqlite3.connect(name)
db.execute("""
CREATE TABLE IF NOT EXISTS states(
id INTEGER PRIMARY KEY AUTOINCREMENT,
value BLOB CHECK (json_valid(value, 0x08))
) STRICT;
""")
yield Storage(db)
db.close()
def fetch_last_state(self) -> StateT:
cursor = self._db.execute(
"""SELECT json(value) FROM states ORDER BY id DESC LIMIT 1;"""
)
rows = list(cursor)
if len(rows) == 0:
return StateT.default()
(value,) = rows[0]
return StateT.model_validate_json(value)
def store_state(self, state: StateT):
c = self._db.cursor()
c.execute("BEGIN;")
try:
c.execute(
"""INSERT INTO states(value) VALUES (jsonb(:value));""",
dict(value=state.model_dump_json())
)
# keep 100 previous versions
c.execute(
"""
DELETE FROM states
WHERE id NOT IN
(SELECT id FROM states ORDER BY id DESC LIMIT 100);
""",
)
c.execute("COMMIT;")
except:
c.execute("ROLLBACK;")
raise