def begin_writer_txn(self,
lmdbenv: lmdb.Environment,
buffer: bool = False) -> lmdb.Transaction:
"""Start a write enabled transaction on the given environment
If multiple write transactions are requested for the same handle, only
one instance of the transaction handle will be returened, and will not
close until all operations on that handle have requested to close
Parameters
----------
lmdbenv : lmdb.Environment
the environment to open the transaction on
buffer : bool, optional
if buffer objects should be used (the default is False, which does
not use buffers)
Returns
-------
lmdb.Transaction
transaction handle to perform operations on
"""
if self.WriterAncestors[lmdbenv] == 0:
self.WriterTxn[lmdbenv] = lmdbenv.begin(write=True, buffers=buffer)
self.WriterAncestors[lmdbenv] += 1
return self.WriterTxn[lmdbenv]
def begin_reader_txn(self,
lmdbenv: lmdb.Environment,
buffer: bool = False) -> lmdb.Transaction:
"""Start a reader only txn for the given environment
If there a read-only transaction for the same environment already exists
then the same reader txn handle will be returned, and will not close
until all operations on that handle have said they are finished.
Parameters
----------
lmdbenv : lmdb.Environment
the environment to start the transaction in.
buffer : bool, optional
weather a buffer transaction should be used (the default is False,
which means no buffers are returned)
Returns
-------
lmdb.Transaction
handle to the lmdb transaction.
"""
if self.ReaderAncestors[lmdbenv] == 0:
self.ReaderTxn[lmdbenv] = lmdbenv.begin(write=False,
buffers=buffer)
self.ReaderAncestors[lmdbenv] += 1
return self.ReaderTxn[lmdbenv]
def retrieve_queue_data(env: lmdb.Environment) -> QueueData:
"""
Retrieve set value in the queue database.
Value stored in the queue database are high water mark values and pruning
strategy.
:param env: environment that stores queue data
:return: data stored in LMDB related to the queue
"""
with env.begin(write=False) as txn:
queue_db = env.open_db(key=QUEUE_DB, txn=txn, create=False)
message_timeout = bytes_to_int(
txn.get(key=MESSAGE_TIMEOUT_KEY, db=queue_db))
max_messages = bytes_to_int(txn.get(key=MAX_MESSAGES_KEY, db=queue_db))
hwm_db_size = bytes_to_int(txn.get(key=HWM_DB_SIZE_KEY, db=queue_db))
strategy = bytes_to_str(txn.get(key=STRATEGY_KEY, db=queue_db))
subscriber_db = env.open_db(key=SUBSCRIBER_DB, txn=txn, create=False)
sub_cursor = txn.cursor(db=subscriber_db)
subscriber_ids = set(
bytes_to_str(subscriber_id)
for subscriber_id in sub_cursor.iternext(values=False))
queue_data = QueueData(message_timeout=message_timeout,
max_messages=max_messages,
hwm_db_size=hwm_db_size,
strategy=strategy,
subscriber_ids=subscriber_ids)
return queue_data
def __init__(self, model: s_datamodel.Model,
dbenv: lmdb.Environment) -> None:
'''
Creates metadata for all the indices.
Args:
dbenv (lmdb.Environment): the lmdb instance in which to store the metadata.
Returns:
None
'''
self._dbenv = dbenv
self.model = model
# The table in the database file (N.B. in LMDB speak, this is called a database)
self._metatbl = dbenv.open_db(b'meta')
is_new_db = False
with dbenv.begin(db=self._metatbl, buffers=True) as txn:
indices_enc = txn.get(b'indices')
progress_enc = txn.get(b'progress')
if indices_enc is None or progress_enc is None:
if indices_enc is None and progress_enc is None:
is_new_db = True
indices_enc = s_msgpack.en({'present': {}, 'deleting': []})
progress_enc = s_msgpack.en({})
else:
raise s_exc.CorruptDatabase(
'missing meta information in index meta'
) # pragma: no cover
indices = s_msgpack.un(indices_enc)
# The details about what the indices are actually indexing: the datapath and type.
self.indices = {
k: _MetaEntry(model, **s_msgpack.un(v))
for k, v in indices.get('present', {}).items()
}
self.deleting = list(indices.get('deleting', ()))
# Keeps track (non-persistently) of which indices have been paused
self.asleep = defaultdict(bool) # type: ignore
# How far each index has progressed as well as statistics
self.progresses = s_msgpack.un(progress_enc)
if not all(p in self.indices for p in self.deleting):
raise s_exc.CorruptDatabase(
'index meta table: deleting entry with unrecognized property name'
) # pragma: no cover
if not all(p in self.indices for p in self.progresses):
raise s_exc.CorruptDatabase(
'index meta table: progress entry with unrecognized property name'
) # pragma: no cover
if is_new_db:
self.persist()
def push(self,
env: lmdb.Environment,
write: bool = False,
iterator: bool = False):
try:
txn = env.begin(
write = write, buffers = True,
parent = self.stacks[env][-1].transaction \
if self.stacks[env] and self.stacks[env][-1].write else None
)
self.stacks[env].append(ExplicitContext(txn, write, iterator))
except lmdb.Error as exc:
raise TransactionError() from exc
def initialize_environment(env: lmdb.Environment, namespace: str) -> str:
try:
txn = None
txn = env.begin(write=True, buffers=False)
env_uuid = txn.get(key=constants.ENVIRONMENT_UUID_KEY.encode('utf-8'))
if env_uuid is None:
env_uuid = str(uuid.uuid4())
assert txn.put(key=constants.ENVIRONMENT_UUID_KEY.encode('utf-8'),
value=env_uuid.encode('utf-8'))
else:
env_uuid = env_uuid.decode('utf-8')
if namespace == constants.ROOT_NAMESPACE:
cursor = txn.cursor()
cursor.first()
while True:
if cursor.key().decode('utf-8') not in [
constants.ENVIRONMENT_UUID_KEY,
constants.ATTRIBUTE_DATABASE,
constants.VERSION_DATABASE,
constants.DESCRIPTOR_DATABASE, constants.NAME_DATABASE
]:
mapsize[cursor.key().decode('utf-8')] = \
struct.unpack('@N', cursor.value())[0]
if not cursor.next():
break
txn.commit()
return env_uuid
except BaseException as exc:
if txn:
try:
txn.abort()
except lmdb.Error:
pass
if isinstance(exc, lmdb.Error):
raise TransactionError() from exc
raise exc
class LmdbDataset:
def __init__(self, databasePath, readonly=True):
self.databasePath = databasePath
self.readonly = readonly
self.__init()
def __init(self):
create = not self.readonly
self.env = Environment(self.databasePath,
map_size=3e9,
subdir=False,
max_dbs=64,
mode=0,
create=create,
readonly=self.readonly)
self.descriptorToDb = {}
self.nextIds = {}
self.__addDb(SceneSetup, create=create)
self.__addDb(ScatterSample, create=create)
self.__addDb(DisneyDescriptor, create=create)
self.__addDb(BakedInterpolationSet, create=create)
self.__addDb(Result, create=create)
def __addDb(self, protocol, create):
name = protocol.DESCRIPTOR.name
db = self.env.open_db(name.encode('ascii'),
integerkey=True,
create=create)
self.descriptorToDb[name] = db
self.nextIds[name] = 0
return db
def append(self, value):
name, db = self.__getNameAndDb(value)
with self.env.begin(write=True) as transaction:
transaction.put((self.nextIds[name]).to_bytes(4, 'little'),
value.SerializeToString(),
db=db)
self.nextIds[name] += 1
def getCountOf(self, protocolType):
name, db = self.__getNameAndDb(protocolType)
with self.env.begin() as transaction:
return transaction.stat(db)['entries']
def get(self, protocolType, id, buffers=False):
_, db = self.__getNameAndDb(protocolType)
with self.env.begin(db=db, buffers=buffers) as transaction:
serialized = transaction.get(id.to_bytes(4, 'little'), db=db)
protocol = protocolType()
protocol.ParseFromString(serialized)
return protocol
def getCountBeforeLastFlatCloud(self):
with self.env.begin() as transaction:
_, db = self.__getNameAndDb(SceneSetup)
cursor = transaction.cursor(db)
cursor.first()
id = 0
for key, value in cursor:
scene = SceneSetup()
scene.ParseFromString(value)
if "RoundClouds" in scene.cloud_path:
return id
id += BATCH_SIZE
def __getNameAndDb(self, protocolType):
name = protocolType.DESCRIPTOR.name
db = self.descriptorToDb[name]
if db == None:
db = self.__addDb(protocolType, False)
return (name, db)
def __getstate__(self):
state = {"databasePath": self.databasePath, "readonly": self.readonly}
return state
def __setstate__(self, state):
# Restore instance attributes
self.__dict__.update(state)
self.__init()