110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
|
# 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)
|