Mongo (#6)

Author: hnipun

labml_db/driver/__init__.py

@@ -7,7 +7,7 @@
 
 
 class DbDriver:
-    def __init__(self, serializer: 'Serializer', model_cls: Type['Model']):
+    def __init__(self, serializer: Optional['Serializer'], model_cls: Type['Model']):
         self.model_name = model_cls.__name__
         self._serializer = serializer
 

labml_db/driver/mongo.py

@@ -0,0 +1,72 @@
+from typing import List, Type, TYPE_CHECKING, Optional, Dict
+
+from labml_db.serializer.utils import encode_keys, decode_keys
+
+from . import DbDriver
+from ..types import ModelDict
+
+if TYPE_CHECKING:
+    import pymongo
+    from ..model import Model
+
+
+class MongoDbDriver(DbDriver):
+    def __init__(self, model_cls: Type['Model'], db: 'pymongo.mongo_client.database.Database'):
+        super().__init__(None, model_cls)
+        self._collection = db[self.model_name]
+
+    def _to_obj_id(self, key: str):
+        return str(key)
+        # return ObjectId(key.split(':')[1])
+
+    def _to_key(self, mongo_id: any):
+        # self.model_name + ':' + str(d['_id']
+        return mongo_id
+
+    def _load_data(self, data: ModelDict):
+        if data is None:
+            return None
+        del data['_id']
+        return decode_keys(data)
+
+    def _dump_data(self, key: str, data: ModelDict):
+        d: Dict = data.copy()
+        d = encode_keys(d)
+        d['_id'] = self._to_obj_id(key)
+
+        return d
+
+    def mload_dict(self, key: List[str]) -> List[Optional[ModelDict]]:
+        obj_keys = [self._to_obj_id(k) for k in key]
+        cursor = self._collection.find({'_id': {'$in': obj_keys}})
+        res = [None for _ in key]
+        idx = {k: i for i, k in enumerate(obj_keys)}
+        for d in cursor:
+            i = idx[d['_id']]
+            res[i] = self._load_data(d)
+
+        return res
+
+    def load_dict(self, key: str) -> Optional[ModelDict]:
+        d = self._collection.find_one({'_id': self._to_obj_id(key)})
+        return self._load_data(d)
+
+    def msave_dict(self, key: List[str], data: List[ModelDict]):
+        objs = [self._dump_data(k, d) for k, d in zip(key, data)]
+
+        from pymongo import ReplaceOne
+        replacements = [ReplaceOne({'_id': d['_id']}, d, True) for d in objs]
+        self._collection.bulk_write(replacements, False)
+
+    def save_dict(self, key: str, data: ModelDict):
+        obj = self._dump_data(key, data)
+
+        self._collection.replace_one({'_id': obj['_id']}, obj, True)
+
+    def delete(self, key: str):
+        self._collection.delete_one({'_id': self._to_obj_id(key)})
+
+    def get_all(self) -> List[str]:
+        cur = self._collection.find(projection=['_id'])
+        keys = [self._to_key(d['_id']) for d in cur]
+        return keys

labml_db/index_driver/__init__.py

@@ -1,4 +1,4 @@
-from typing import Type, TYPE_CHECKING, List
+from typing import Type, TYPE_CHECKING, List, Optional
 
 if TYPE_CHECKING:
     from . import Index
@@ -11,10 +11,10 @@ def __init__(self, index_cls: Type['Index']):
     def delete(self, index_key: str):
         raise NotImplementedError
 
-    def get(self, index_key: str) -> str:
+    def get(self, index_key: str) -> Optional[str]:
         raise NotImplementedError
 
-    def mget(self, index_key: List[str]) -> List[str]:
+    def mget(self, index_key: List[str]) -> List[Optional[str]]:
         raise NotImplementedError
 
     def set(self, index_key: str, model_key: str):

labml_db/index_driver/mongo.py

@@ -0,0 +1,41 @@
+from typing import Dict, Type, Optional, TYPE_CHECKING, List
+
+from . import IndexDbDriver
+
+if TYPE_CHECKING:
+    import pymongo
+    from .. import Index
+
+
+class MongoIndexDbDriver(IndexDbDriver):
+    _cache: Optional[Dict[str, str]]
+
+    def __init__(self, index_cls: Type['Index'], db: 'pymongo.mongo_client.database.Database'):
+        super().__init__(index_cls)
+        self._index = db[f'_index_{self.index_name}']
+
+    def delete(self, index_key: str):
+        self._index.delete_one({'_id': index_key})
+
+    def get(self, index_key: str) -> Optional[str]:
+        d = self._index.find_one({'_id': index_key})
+        if d is None:
+            return None
+        return d['value']
+
+    def mget(self, index_key: List[str]) -> List[Optional[str]]:
+        cursor = self._index.find({'_id': {'$in': index_key}})
+        res = [None for _ in index_key]
+        idx = {k: i for i, k in enumerate(index_key)}
+        for d in cursor:
+            i = idx[d['_id']]
+            res[i] = d['value']
+
+        return res
+
+    def set(self, index_key: str, model_key: str):
+        self._index.replace_one({'_id': index_key}, {'_id': index_key, 'value': model_key}, True)
+
+    def get_all(self):
+        cur = self._index.find()
+        return [d['value'] for d in cur]

setup.py

@@ -5,15 +5,15 @@
 
 setuptools.setup(
     name='labml_db',
-    version='0.0.14',
+    version='0.0.15',
     author="Varuna Jayasiri, Nipun Wijerathne",
     author_email="[email protected], [email protected]",
-    description="Minimalistic ORM for JSON/YAML/Pickle file based DB",
+    description="Minimalistic ORM for JSON/YAML/Pickle file based/redis/mongo DB",
     long_description=long_description,
     long_description_content_type="text/x-rst",
-    url="https://github.com/lab-ml/db",
+    url="https://github.com/labmlai/db",
     project_urls={
-        'Documentation': 'https://lab-ml.com/'
+        'Documentation': 'https://labml.ai'
     },
     packages=setuptools.find_packages(exclude=('test',
                                                'test.*')),

test/mongo/__init__.py

@@ -0,0 +1,43 @@
+from labml import monit
+from labml.logger import inspect
+from pymongo import MongoClient
+
+CONNECTION_STRING = "mongodb://localhost:27017/"
+
+
+def get_database():
+    client = MongoClient(CONNECTION_STRING)
+
+    inspect(client)
+    # inspect(client.server_info())
+    inspect(client.list_database_names())
+
+    # client.drop_database('papers')
+
+    return client['papers']
+
+
+def test():
+    db = get_database()
+    inspect(db)
+    collection = db['papers']
+    inspect(collection)
+
+    with monit.section('Insert'):
+        paper_1 = {
+            "_id": "abcd",
+            "title": "Attention is all you need",
+            "abstract": "abstract"
+        }
+
+        collection.insert_many([paper_1])
+
+    with monit.section('Find'):
+        item_details = collection.find()
+        for item in item_details:
+            inspect(type(item))
+            inspect(item)
+
+
+if __name__ == "__main__":
+    test()

test/mongo/api_check.py

@@ -0,0 +1,57 @@
+from pymongo import MongoClient
+
+from labml_db import Model, Index
+from labml_db.driver.mongo import MongoDbDriver
+from labml_db.index_driver.mongo import MongoIndexDbDriver
+from test.simple import User, Project, UsernameIndex
+
+
+def test_setup():
+    db = MongoClient("mongodb://localhost:27017/")['testdb']
+    Model.set_db_drivers([
+        MongoDbDriver(User, db),
+        MongoDbDriver(Project, db)
+    ])
+    Index.set_db_drivers([
+        MongoIndexDbDriver(UsernameIndex, db)
+    ])
+
+
+def test():
+    proj = Project(name='nlp')
+    user = User(name='John')
+    user.projects.append(proj.key)
+    user.occupation = 'test'
+    user2 = User(name='X')
+    print(user.projects, user2.projects)
+    user.save()
+    proj.save()
+
+    print(user.projects[0].load().name)
+
+
+def test_load():
+    keys = User.get_all()
+    print([k.load() for k in keys])
+    keys = Project.get_all()
+    print([k.load().name for k in keys])
+
+
+def test_index():
+    user_key = UsernameIndex.get('John')
+    if user_key:
+        print('index', user_key.load())
+        user_key.delete()
+
+    user = User(name='John')
+    user.save()
+    UsernameIndex.set(user.name, user.key)
+
+    print(user.key, user.name, user.projects)
+
+
+if __name__ == '__main__':
+    test_setup()
+    test()
+    test_load()
+    test_index()