hat.event.backends.lmdb.environment

  1from pathlib import Path
  2import asyncio
  3import typing
  4
  5import lmdb
  6
  7from hat import aio
  8
  9from hat.event.backends.lmdb import common
 10
 11
 12async def create(db_path: Path,
 13                 max_size: int = common.default_max_size
 14                 ) -> 'Environment':
 15    env = Environment()
 16    env._loop = asyncio.get_running_loop()
 17    env._async_group = aio.Group()
 18    env._executor = aio.Executor(1, log_exceptions=False)
 19    env._env = None
 20    env._dbs = {}
 21
 22    env.async_group.spawn(aio.call_on_cancel, env._on_close)
 23
 24    try:
 25        await env._executor.spawn(env._ext_init, db_path, max_size)
 26
 27    except BaseException:
 28        await aio.uncancellable(env.async_close())
 29        raise
 30
 31    return env
 32
 33
 34class Environment(aio.Resource):
 35
 36    @property
 37    def async_group(self) -> aio.Group:
 38        return self._async_group
 39
 40    async def execute(self,
 41                      ext_fn: typing.Callable,
 42                      *args: typing.Any) -> typing.Any:
 43        if not self.is_open:
 44            raise Exception('environment closed')
 45
 46        return await self._executor.spawn(ext_fn, *args)
 47
 48    def ext_begin(self, write: bool = False) -> lmdb.Transaction:
 49        return self._env.begin(write=write,
 50                               buffers=True)
 51
 52    def ext_cursor(self,
 53                   txn: lmdb.Transaction,
 54                   db_type: common.DbType
 55                   ) -> lmdb.Cursor:
 56        return txn.cursor(self._dbs[db_type])
 57
 58    def ext_stat(self,
 59                 txn: lmdb.Transaction,
 60                 db_type: common.DbType
 61                 ) -> dict[str, int]:
 62        return txn.stat(self._dbs[db_type])
 63
 64    def ext_read(self,
 65                 txn: lmdb.Transaction,
 66                 db_type: common.DbType
 67                 ) -> typing.Iterable[tuple[common.DbKey, common.DbValue]]:
 68        db_def = common.db_defs[db_type]
 69        with txn.cursor(self._dbs[db_type]) as cursor:
 70            for encoded_key, encoded_value in cursor:
 71                key = db_def.decode_key(encoded_key)
 72                value = db_def.decode_value(encoded_value)
 73                yield key, value
 74
 75    def ext_write(self,
 76                  txn: lmdb.Transaction,
 77                  db_type: common.DbType,
 78                  data: typing.Iterable[tuple[common.DbKey, common.DbValue]]):
 79        db_def = common.db_defs[db_type]
 80        with txn.cursor(self._dbs[db_type]) as cursor:
 81            for key, value in data:
 82                encoded_key = db_def.encode_key(key)
 83                encoded_value = db_def.encode_value(value)
 84                cursor.put(encoded_key, encoded_value)
 85
 86    async def _on_close(self):
 87        if self._env:
 88            await self._executor.spawn(self._env.close)
 89
 90        await self._executor.async_close()
 91
 92    def _ext_init(self, db_path, max_size):
 93        create = not db_path.exists()
 94
 95        self._env = common.ext_create_env(path=db_path,
 96                                          max_size=max_size)
 97        self._dbs = {db_type: common.ext_open_db(env=self._env,
 98                                                 db_type=db_type,
 99                                                 create=create)
100                     for db_type in common.DbType}