You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
261 lines
6.9 KiB
261 lines
6.9 KiB
|
|
import os |
|
import lmdb |
|
import shutil |
|
from itertools import islice |
|
|
|
from ..types import unicode_, to_bytes, bson |
|
from . import ( |
|
AbstractStorage, |
|
AbstractDatabase, |
|
AbstractCollection, |
|
AbstractCursor, |
|
|
|
StorageDuplicateKeyError, |
|
) |
|
|
|
|
|
LMDB_DB_EXT = ".mdb" |
|
|
|
|
|
class LMDBKVEngine(object): |
|
"""Per collection""" |
|
|
|
dbname = to_bytes("documents") |
|
|
|
def __init__(self, options=None): |
|
""" |
|
""" |
|
opt = options or dict() |
|
opt.update({ |
|
"subdir": False, |
|
"max_dbs": 1, |
|
}) |
|
self.opt = opt |
|
self._path = None |
|
|
|
def set_path(self, path): |
|
self._path = path |
|
|
|
def open(self, readonly=None): |
|
opt = self.opt.copy() |
|
if readonly is not None: |
|
opt["readonly"] = readonly |
|
|
|
return lmdb.open(self._path, **opt) |
|
|
|
def iter_docs(self): |
|
if not os.path.isfile(self._path): |
|
return |
|
|
|
with self.open(readonly=True) as env: |
|
db = env.open_db(self.dbname) |
|
with env.begin(db, write=False) as txn: |
|
cursor = txn.cursor() |
|
for encoded_doc in cursor.iternext(keys=False, values=True): |
|
yield encoded_doc |
|
|
|
def write(self, environment, pairs, overwrite=False): |
|
if not os.path.isfile(self._path): |
|
return |
|
|
|
dup = False |
|
with environment as env: |
|
db = env.open_db(self.dbname) |
|
with env.begin(db, write=True) as txn: |
|
for doc_id, encoded_doc in pairs: |
|
id = bson.id_encode(doc_id) |
|
if not txn.put(id, encoded_doc, overwrite=overwrite): |
|
dup = True |
|
break |
|
if dup: |
|
raise StorageDuplicateKeyError() |
|
|
|
def delete(self, environment, doc_ids): |
|
if not os.path.isfile(self._path): |
|
return |
|
|
|
with environment as env: |
|
db = env.open_db(self.dbname, reverse_key=False) |
|
with env.begin(db, write=True) as txn: |
|
cursor = txn.cursor() |
|
for doc_id in doc_ids: |
|
id = bson.id_encode(doc_id) |
|
if cursor.set_key(id): |
|
cursor.delete() |
|
|
|
|
|
class LMDBStorage(AbstractStorage): |
|
""" |
|
""" |
|
|
|
def __init__(self, repository, storage_config): |
|
super(LMDBStorage, self).__init__(repository, storage_config) |
|
self._conn = LMDBKVEngine(self._config) |
|
|
|
def _db_path(self, db_name): |
|
""" |
|
Get Monty database dir path. |
|
""" |
|
return os.path.join(self._repository, db_name) |
|
|
|
@classmethod |
|
def nice_name(cls): |
|
return "lightning" |
|
|
|
@classmethod |
|
def config(cls, map_size=10485760, **kwargs): |
|
""" |
|
""" |
|
return { |
|
"map_size": int(map_size), |
|
} |
|
|
|
def database_create(self, db_name): |
|
if not os.path.isdir(self._db_path(db_name)): |
|
os.makedirs(self._db_path(db_name)) |
|
|
|
def database_drop(self, db_name): |
|
db_path = self._db_path(db_name) |
|
if os.path.isdir(db_path): |
|
shutil.rmtree(db_path) |
|
|
|
def database_list(self): |
|
return [ |
|
name for name in os.listdir(unicode_(self._repository)) |
|
if os.path.isdir(self._db_path(name)) |
|
] |
|
|
|
|
|
class LMDBDatabase(AbstractDatabase): |
|
""" |
|
""" |
|
|
|
def __init__(self, storage, subject): |
|
super(LMDBDatabase, self).__init__(storage, subject) |
|
self._db_path = storage._db_path(self._name) |
|
self._conn = storage._conn |
|
|
|
def _col_path(self, col_name): |
|
""" |
|
Get database file path, which is Monty collection. |
|
""" |
|
return os.path.join(self._db_path, col_name) + LMDB_DB_EXT |
|
|
|
def database_exists(self): |
|
return os.path.isdir(self._db_path) |
|
|
|
def collection_exists(self, col_name): |
|
return os.path.isfile(self._col_path(col_name)) |
|
|
|
def collection_create(self, col_name): |
|
if not self.database_exists(): |
|
self._storage.database_create(self._name) |
|
self._conn.set_path(self._col_path(col_name)) |
|
environment = self._conn.open() |
|
|
|
return environment |
|
|
|
def collection_drop(self, col_name): |
|
if self.collection_exists(col_name): |
|
os.remove(self._col_path(col_name)) |
|
|
|
def collection_list(self): |
|
if not self.database_exists(): |
|
return [] |
|
return [os.path.splitext(name)[0] |
|
for name in os.listdir(unicode_(self._db_path)) |
|
if name.endswith(LMDB_DB_EXT)] |
|
|
|
|
|
LMDBStorage.contractor_cls = LMDBDatabase |
|
|
|
|
|
class LMDBCollection(AbstractCollection): |
|
""" |
|
""" |
|
|
|
def __init__(self, database, subject): |
|
super(LMDBCollection, self).__init__(database, subject) |
|
self._conn = database._conn |
|
self._conn.set_path(database._col_path(self._name)) |
|
|
|
def _ensure_table(func): |
|
def make_table(self, *args, **kwargs): |
|
environment = None |
|
if not self._database.collection_exists(self._name): |
|
environment = self._database.collection_create(self._name) |
|
return func(self, env=environment, *args, **kwargs) |
|
return make_table |
|
|
|
@_ensure_table |
|
def write_one(self, doc, check_keys=True, env=None): |
|
env = env or self._conn.open() |
|
|
|
id = doc["_id"] |
|
encoded = self._encode_doc(doc, check_keys) |
|
self._conn.write(env, [(id, encoded)]) |
|
|
|
return id |
|
|
|
@_ensure_table |
|
def write_many(self, docs, check_keys=True, ordered=True, env=None): |
|
env = env or self._conn.open() |
|
ids = list() |
|
|
|
def produce_encoded_docs(): |
|
for doc in docs: |
|
id = doc["_id"] |
|
yield id, self._encode_doc(doc, check_keys) |
|
ids.append(id) |
|
|
|
self._conn.write(env, produce_encoded_docs()) |
|
|
|
return ids |
|
|
|
def update_one(self, doc): |
|
env = self._conn.open() |
|
|
|
id = doc["_id"] |
|
encoded = self._encode_doc(doc) |
|
self._conn.write(env, [(id, encoded)], overwrite=True) |
|
|
|
def update_many(self, docs): |
|
env = self._conn.open() |
|
|
|
def produce_encoded_docs(): |
|
for doc in docs: |
|
yield doc["_id"], self._encode_doc(doc) |
|
|
|
self._conn.write(env, produce_encoded_docs(), overwrite=True) |
|
|
|
def delete_one(self, id): |
|
env = self._conn.open() |
|
self._conn.delete(env, [id]) |
|
|
|
def delete_many(self, ids): |
|
env = self._conn.open() |
|
self._conn.delete(env, ids) |
|
|
|
|
|
LMDBDatabase.contractor_cls = LMDBCollection |
|
|
|
|
|
class LMDBCursor(AbstractCursor): |
|
""" |
|
""" |
|
|
|
def __init__(self, collection, subject): |
|
super(LMDBCursor, self).__init__(collection, subject) |
|
self._conn = self._collection._conn |
|
|
|
def query(self, max_scan): |
|
docs = (self._decode_doc(doc) for doc in self._conn.iter_docs()) |
|
|
|
if not max_scan: |
|
return docs |
|
else: |
|
return islice(docs, max_scan) |
|
|
|
|
|
LMDBCollection.contractor_cls = LMDBCursor
|
|
|