update

devplacepy_services/base/sqlite_broker.py

@@ -0,0 +1,110 @@
+# retoor <retoor@molodetz.nl>
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, TypeVar
+
+import dataset
+
+from devplacepy_services.base.config import sqlite_read_pool_size
+
+T = TypeVar("T")
+
+PRAGMAS = [
+    "PRAGMA journal_mode=WAL",
+    "PRAGMA synchronous=NORMAL",
+    "PRAGMA busy_timeout=30000",
+    "PRAGMA cache_size=-8000",
+    "PRAGMA temp_store=MEMORY",
+    "PRAGMA mmap_size=268435456",
+]
+
+
+class SQLiteFileBroker:
+    name: str
+    path: Path
+    read_pool_size: int
+
+    def __init__(self, name: str, path: Path, read_pool_size: int | None = None) -> None:
+        self.name = name
+        self.path = Path(path)
+        self.read_pool_size = read_pool_size or sqlite_read_pool_size()
+        self._write_db: Any = None
+        self._read_pool: asyncio.Queue[Any] | None = None
+        self._write_queue: asyncio.Queue[Any] | None = None
+        self._write_worker_task: asyncio.Task | None = None
+
+    def _connect_rw(self):
+        return dataset.connect(
+            f"sqlite:///{self.path}",
+            engine_kwargs={
+                "connect_args": {
+                    "timeout": 30,
+                    "check_same_thread": False,
+                },
+            },
+            on_connect_statements=PRAGMAS,
+        )
+
+    def _connect_ro(self):
+        abs_path = self.path.resolve().as_posix()
+        return dataset.connect(
+            f"sqlite:///file:{abs_path}?uri=true&mode=ro",
+            engine_kwargs={
+                "connect_args": {
+                    "timeout": 30,
+                    "check_same_thread": False,
+                    "uri": True,
+                },
+            },
+            on_connect_statements=PRAGMAS,
+        )
+
+    async def startup(self) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._write_db = self._connect_rw()
+        self._read_pool = asyncio.Queue()
+        for _ in range(self.read_pool_size):
+            await self._read_pool.put(self._connect_ro())
+        self._write_queue = asyncio.Queue()
+        self._write_worker_task = asyncio.create_task(self._write_worker())
+
+    async def shutdown(self) -> None:
+        if self._write_worker_task is not None:
+            self._write_worker_task.cancel()
+            try:
+                await self._write_worker_task
+            except asyncio.CancelledError:
+                pass
+            self._write_worker_task = None
+
+    async def read(self, fn: Callable[..., T], *args, **kwargs) -> T:
+        assert self._read_pool is not None
+        db = await self._read_pool.get()
+        try:
+            return await asyncio.to_thread(fn, db, *args, **kwargs)
+        finally:
+            await self._read_pool.put(db)
+
+    async def write(self, fn: Callable[..., T], *args, **kwargs) -> T:
+        assert self._write_queue is not None
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+        await self._write_queue.put((fn, args, kwargs, future))
+        return await future
+
+    async def _write_worker(self) -> None:
+        assert self._write_queue is not None
+        while True:
+            fn, args, kwargs, future = await self._write_queue.get()
+            try:
+                result = fn(self._write_db, *args, **kwargs)
+                future.set_result(result)
+            except Exception as exc:
+                future.set_exception(exc)
+
+    async def read_after_write(self, fn: Callable[..., T], *args, **kwargs) -> T:
+        return await asyncio.to_thread(fn, self._write_db, *args, **kwargs)