diff --git a/.gitignore b/.gitignore index 0b96173..c00be02 100644 --- a/.gitignore +++ b/.gitignore @@ -223,3 +223,4 @@ cython_debug/ # Other /dev.py /*.db +/*.db.* diff --git a/ptvp35/__init__.py b/ptvp35/__init__.py index 6aad4d4..cd8cc3b 100644 --- a/ptvp35/__init__.py +++ b/ptvp35/__init__.py @@ -3,23 +3,27 @@ import json import pathlib import pickle import shutil +import traceback from io import StringIO from typing import Any, Optional, IO, Type, Hashable class Request: + def __init__(self, future: Optional[asyncio.Future]): + self.future = future + def set_result(self, result): - raise NotImplementedError + self.future and self.future.set_result(result) def set_exception(self, exception): - raise NotImplementedError + self.future and self.future.set_exception(exception) class KVRequest(Request): def __init__(self, key: Any, value: Any, future: Optional[asyncio.Future]): + super().__init__(future) self.key = key self.value = value - self.future = future def free(self): return type(self)(self.key, self.value, None) @@ -31,22 +35,9 @@ class KVRequest(Request): def fromline(cls, line: str) -> 'KVRequest': raise NotImplementedError - def set_result(self, result): - self.future and self.future.set_result(result) - - def set_exception(self, exception): - self.future and self.future.set_exception(exception) - class DumpRequest(Request): - def __init__(self, future: Optional[asyncio.Future]): - self.future = future - - def set_result(self, result): - self.future and self.future.set_result(result) - - def set_exception(self, exception): - self.future and self.future.set_exception(exception) + pass class UnkownRequestType(TypeError): @@ -83,10 +74,19 @@ class KVJson(KVRequest): return KVJson(cls._load_key(d['key']), d['value'], None) +class ErrorRequest(Request): + def __init__(self, line: str, future: Optional[asyncio.Future]): + super().__init__(future) + self.line = line + + async def wait(self): + await self.future + + class Db: __mmdb: Optional[dict] __loop: Optional[asyncio.AbstractEventLoop] - __queue: Optional[asyncio.Queue] + __queue: Optional[asyncio.Queue[Request]] __file: Optional[IO[str]] __buffer: Optional[StringIO] __task: Optional[asyncio.Future] @@ -98,14 +98,24 @@ class Db: self.__path = pathlib.Path(path) self.__path_backup = pathlib.Path(path + '.backup') self.__path_recover = pathlib.Path(path + '.recover') + self.__path_error = pathlib.Path(path + '.error') self.__task = None + def _queue_error(self, line: str): + request = ErrorRequest(line, self.__loop.create_future()) + self.__queue.put_nowait(request) + self.__loop.create_task(request.wait()) + def io2db(self, io: IO[str], db: dict) -> int: size = 0 for line in io: - request = self.kvrequest_type.fromline(line) - db[request.key] = request.value - size += len(line) + try: + request = self.kvrequest_type.fromline(line) + db[request.key] = request.value + size += len(line) + except (json.JSONDecodeError, pickle.UnpicklingError, EOFError): + traceback.print_exc() + self._queue_error(line) return size def db2io(self, db: dict, io: IO[str]) -> int: @@ -166,12 +176,18 @@ class Db: await self._do_dump_buffer() await self._reload_if_oversized() + async def _save_error(self, line: str): + with open(self.__path_error, 'a') as file: + await self.__loop.run_in_executor(None, file.write, line.strip() + '\n') + async def _handle_request(self, request: Request): if isinstance(request, self.kvrequest_type): await self._write(request.line(), request) elif isinstance(request, DumpRequest): await self._dump_buffer() request.set_result(None) + elif isinstance(request, ErrorRequest): + await self._save_error(request.line) else: raise UnkownRequestType @@ -215,18 +231,21 @@ class Db: async def _initialize_mmdb(self): self.__mmdb = {} await self._load_from_file() + self.__file = open(self.__path, "a") async def _initialize_queue(self): self.__queue = asyncio.Queue() - self.__file = open(self.__path, "a") self.__buffer = StringIO() + + async def _start_task(self): self.__task = self.__loop.create_task(self._background_task()) async def _initialize(self): assert self.__task is None self.__loop = asyncio.get_event_loop() - await self._initialize_mmdb() await self._initialize_queue() + await self._initialize_mmdb() + await self._start_task() async def __aenter__(self): await self._initialize()