+import os
+import sqlite3
+import sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+from tempfile import gettempdir, NamedTemporaryFile
+
+
+import aiosqlite
+
+
+def get_tempdir() -> Path:
+ tmpdir = Path("/dev/shm")
+ if not tmpdir.exists():
+ tmpdir = Path(gettempdir())
+ assert tmpdir.exists()
+ return tmpdir
+
+
+class ActivityManager:
+ tmp_ctx: NamedTemporaryFile
+ ENVIRON_KEY = "_ACTIVITY_CTX"
+ worker_id: int
+
+ @classmethod
+ def setup(cls):
+ cls.tmp_ctx = NamedTemporaryFile(
+ suffix=".db", prefix="activity", dir=get_tempdir(), delete_on_close=False
+ )
+ cls.tmp_ctx.__enter__().close()
+ os.environ[cls.ENVIRON_KEY] = tmp_path = cls.tmp_ctx.name
+ with sqlite3.connect(tmp_path) as db:
+ db.execute("PRAGMA journal_mode = WAL")
+ db.executescript(
+ """
+ CREATE TABLE worker (id INTEGER PRIMARY KEY, pid INTEGER) STRICT;
+ CREATE TABLE websocket (
+ id INTEGER PRIMARY KEY,
+ worker_id INTEGER,
+ user_id INTEGER,
+ FOREIGN KEY(worker_id) REFERENCES worker(id)
+ ) STRICT;
+ CREATE TABLE user_status (
+ id INTEGER PRIMARY KEY,
+ user_id INTEGER UNIQUE,
+ status TEXT
+ ) STRICT;
+ """
+ )
+ db.commit()
+
+ @classmethod
+ def cleanup(cls):
+ tmp_path = cls.tmp_ctx.name
+ for suffix in ("-shm", "-wal"):
+ Path(f"{tmp_path}{suffix}").unlink(missing_ok=True)
+ cls.tmp_ctx.__exit__(*sys.exc_info())
+
+ def __init__(self):
+ self.db = None
+
+ @asynccontextmanager
+ async def db_block(self):
+ try:
+ yield
+ await self.db.commit()
+ except Exception as original_error:
+ try:
+ await self.db.rollback()
+ except Exception as rollback_error:
+ raise ExceptionGroup(
+ "Transaction and rollback both failed",
+ [original_error, rollback_error],
+ ) from original_error
+ raise
+
+ async def insert_returning_id(self, sql, params):
+ async with await self.db.execute(sql, params) as cursor:
+ return cursor.lastrowid
+
+ async def execute_fetchall(self, sql, params, cursor=None):
+ if cursor is None:
+ async with await self.db.execute(sql, params) as cursor:
+ return await cursor.fetchall()
+ await cursor.execute(sql, params)
+ return await cursor.fetchall()
+
+ async def __aenter__(self):
+ self.db = await aiosqlite.connect(os.environ[self.ENVIRON_KEY])
+ self.db.row_factory = aiosqlite.Row
+ await self.db.execute("PRAGMA synchronous=NORMAL")
+ await self.db.execute("PRAGMA foreign_keys=ON")
+ await self.db.execute("PRAGMA busy_timeout=5000")
+ async with self.db_block():
+ self.worker_id = await self.insert_returning_id(
+ "INSERT INTO worker (pid) VALUES (?)", (os.getpid(),)
+ )
+ return self
+
+ async def register_websocket(self, user_id=0):
+ async with self.db_block():
+ websocket_id = await self.insert_returning_id(
+ "INSERT INTO websocket (worker_id, user_id) VALUES (?, ?)",
+ (self.worker_id, user_id),
+ )
+ if user_id:
+ async with await self.db.cursor() as cursor:
+ await cursor.execute(
+ (
+ "INSERT OR IGNORE INTO user_status (user_id, status) "
+ "VALUES (?, ?)"
+ ),
+ (user_id, "online"),
+ )
+ return websocket_id
+
+ async def update_websocket_status(self, websocket_id, new_status):
+ async with self.db_block(), await self.db.cursor() as cursor:
+ await cursor.execute(
+ (
+ "UPDATE user_status SET status=? WHERE user_id=("
+ " SELECT user_id FROM websocket WHERE worker_id=? AND id=?"
+ ")"
+ ),
+ (new_status, self.worker_id, websocket_id),
+ )
+
+ async def get_websocket_status(self, user_ids):
+ p = ", ".join("?" for _ in range(len(user_ids)))
+ return await self.execute_fetchall(
+ f"SELECT user_id, status FROM user_status WHERE user_id IN ({p})",
+ user_ids,
+ )
+
+ async def cleanup_user(self, cursor, user_id):
+ if not await self.execute_fetchall(
+ "SELECT id FROM websocket WHERE user_id=?",
+ (user_id,),
+ cursor,
+ ):
+ await cursor.execute(
+ "DELETE FROM user_status WHERE user_id=?",
+ (user_id,)
+ )
+
+ async def unregister_websocket(self, websocket_id, user_id):
+ async with self.db_block(), await self.db.cursor() as cursor:
+ await cursor.execute("DELETE FROM websocket WHERE id=?", (websocket_id,))
+ if user_id:
+ await self.cleanup_user(cursor, user_id)
+
+ async def __aexit__(self, et, exc, tb):
+ try:
+ if not self.worker_id:
+ return
+ async with self.db_block(), await self.db.cursor() as cursor:
+ for row in await self.execute_fetchall(
+ "SELECT user_id FROM websocket WHERE worker_id=?",
+ (self.worker_id,),
+ cursor,
+ ):
+ await self.cleanup_user(cursor, row[0])
+ await cursor.execute(
+ "DELETE FROM websocket WHERE worker_id=?", (self.worker_id,)
+ )
+ await cursor.execute(
+ "DELETE FROM worker WHERE id=?", (self.worker_id,)
+ )
+ finally:
+ await self.db.close()
+ self.db = None
projects / chat / commitdiff