Merge pull request #101 from poissoncorp/RDBC-471-5.0

RDBC-471 Add support for storing @DataClass() objects (for 5.0)

Author: gregolsky

pyravendb/store/document_session.py

@@ -10,6 +10,7 @@
 from .session_timeseries import TimeSeries
 from .session_counters import DocumentCounters
 from typing import Dict, List
+from collections import MutableSet
 
 
 class _SaveChangesData(object):
@@ -19,6 +20,129 @@ def __init__(self, commands, deferred_command_count, entities=None):
         self.deferred_command_count = deferred_command_count
 
 
+class _RefEq:
+    def __init__(self, ref):
+        if isinstance(ref, _RefEq):
+            self.ref = ref.ref
+            return
+        self.ref = ref
+
+    # As we split the hashable and unhashable items into separate collections, we only compare _RefEq to other _RefEq
+    def __eq__(self, other):
+        if isinstance(other, _RefEq):
+            return id(self.ref) == id(other.ref)
+        raise TypeError("Expected _RefEq type object")
+
+    def __hash__(self):
+        return id(self.ref)
+
+
+class _RefEqEntityHolder(object):
+    def __init__(self):
+        self.unhashable_items = dict()
+
+    def __len__(self):
+        return len(self.unhashable_items)
+
+    def __contains__(self, item):
+        return _RefEq(item) in self.unhashable_items
+
+    def __delitem__(self, key):
+        del self.unhashable_items[_RefEq(key)]
+
+    def __setitem__(self, key, value):
+        self.unhashable_items[_RefEq(key)] = value
+
+    def __getitem__(self, key):
+        return self.unhashable_items[_RefEq(key)]
+
+    def __getattribute__(self, item):
+        if item == "unhashable_items":
+            return super().__getattribute__(item)
+        return self.unhashable_items.__getattribute__(item)
+
+
+class _DocumentsByEntityHolder(object):
+    def __init__(self):
+        self._hashable_items = dict()
+        self._unhashable_items = _RefEqEntityHolder()
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}: {[item for item in self.__iter__()]}"
+
+    def __len__(self):
+        return len(self._hashable_items) + len(self._unhashable_items)
+
+    def __contains__(self, item):
+        try:
+            return item in self._hashable_items
+        except TypeError as e:
+            if str(e.args[0]).startswith("unhashable type"):
+                return item in self._unhashable_items
+            raise e
+
+    def __setitem__(self, key, value):
+        try:
+            self._hashable_items[key] = value
+        except TypeError as e:
+            if str(e.args[0]).startswith("unhashable type"):
+                self._unhashable_items[key] = value
+                return
+            raise e
+
+    def __getitem__(self, key):
+        try:
+            return self._hashable_items[key]
+        except (TypeError, KeyError):
+            return self._unhashable_items[key]
+
+    def __iter__(self):
+        d = list(map(lambda x: x.ref, self._unhashable_items.keys()))
+        if len(self._hashable_items) > 0:
+            d.extend(self._hashable_items.keys())
+        return (item for item in d)
+
+    def get(self, key, default=None):
+        return self[key] if key in self else default
+
+    def pop(self, key, default_value=None):
+        result = self._hashable_items.pop(key, None)
+        if result is not None:
+            return result
+        return self._unhashable_items.pop(_RefEq(key), default_value)
+
+    def clear(self):
+        self._hashable_items.clear()
+        self._unhashable_items.clear()
+
+
+class _DeletedEntitiesHolder(MutableSet):
+    def __init__(self, items=None):
+        if items is None:
+            items = []
+        self.items = set(map(_RefEq, items))
+
+    def __getattribute__(self, item):
+        if item in ["add", "discard", "items"]:
+            return super().__getattribute__(item)
+        return self.items.__getattribute__(item)
+
+    def __contains__(self, item: object) -> bool:
+        return _RefEq(item) in self.items
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __iter__(self):
+        return (item.ref for item in self.items)
+
+    def add(self, element: object) -> None:
+        return self.items.add(_RefEq(element))
+
+    def discard(self, element: object) -> None:
+        return self.items.discard(_RefEq(element))
+
+
 class DocumentSession(object):
     def __init__(self, database, document_store, requests_executor, session_id, **kwargs):
         """
@@ -33,8 +157,8 @@ def __init__(self, database, document_store, requests_executor, session_id, **kw
         self._requests_executor = requests_executor
         self._documents_by_id = {}
         self._included_documents_by_id = {}
-        self._deleted_entities = set()
-        self._documents_by_entity = {}
+        self._deleted_entities = _DeletedEntitiesHolder()
+        self._documents_by_entity = _DocumentsByEntityHolder()
         self._timeseries_defer_commands = {}
         self._time_series_by_document_id = {}
         self._counters_defer_commands = {}

pyravendb/tests/session_tests/test_delete.py

@@ -3,11 +3,19 @@
 
 sys.path.append(os.path.abspath(__file__ + "/../"))
 
+from dataclasses import dataclass
 from pyravendb.tests.test_base import TestBase
 from pyravendb.custom_exceptions import exceptions
 import unittest
 
 
+@dataclass()
+class Fish:
+    name: str
+    weight: int
+    Id: str
+
+
 class Product(object):
     def __init__(self, Id=None, name=None):
         self.Id = Id
@@ -68,6 +76,19 @@ def test_delete_with_entity(self):
 
             self.assertIsNone(session.load(product.Id))
 
+    def test_delete_dataclass(self):
+        fishie = Fish(name="Tuna", weight=100, Id=None)
+        with self.store.open_session() as session:
+            session.store(fishie)
+            session.save_changes()
+
+        with self.store.open_session() as session:
+            fishie = session.load(fishie.Id)
+            session.delete(fishie)
+            session.save_changes()
+
+            self.assertIsNone(session.load(fishie.Id))
+
 
 if __name__ == "__main__":
     unittest.main()

pyravendb/tests/session_tests/test_documents_by_entity.py

@@ -0,0 +1,23 @@
+from pyravendb.tests.test_base import TestBase
+from pyravendb.store.document_session import _RefEq
+from dataclasses import dataclass
+
+
+@dataclass()
+class TestData:
+    entry: str
+
+
+class TestDocumentsByEntity(TestBase):
+    def setUp(self):
+        super(TestDocumentsByEntity, self).setUp()
+
+    def test_ref_eq(self):
+        data = TestData(entry="classified")
+
+        wrapped_data_1 = _RefEq(data)
+        wrapped_data_2 = _RefEq(wrapped_data_1)
+
+        self.assertTrue(wrapped_data_2 == wrapped_data_1)
+        with self.assertRaises(TypeError):
+            wrapped_data_1 == data

pyravendb/tests/session_tests/test_load.py

@@ -1,9 +1,24 @@
 from pyravendb.commands.raven_commands import PutDocumentCommand
 from pyravendb.tests.test_base import TestBase
 from pyravendb.custom_exceptions import exceptions
+from dataclasses import dataclass
 import unittest
 
 
+@dataclass()
+class OrderData(object):
+    Id: str
+    name: str
+    key: str
+    product_id: str
+
+
+@dataclass()
+class ProductData(object):
+    Id: str
+    name: str
+
+
 class Product(object):
     def __init__(self, Id=None, name=""):
         self.Id = Id
@@ -116,6 +131,18 @@ def test_load_with_include(self):
             session.load("products/101")
         self.assertEqual(session.number_of_requests_in_session, 1)
 
+    def test_load_with_include_dataclass(self):
+        with self.store.open_session() as session:
+            session.store(OrderData("orderd/1", "some_order", "some_key", "productd/1"))
+            session.store(ProductData("productd/1", "some_product"))
+            session.save_changes()
+
+        with self.store.open_session() as session:
+            session.load("orderd/1", includes="product_id")
+            product = session.load("productd/1")
+        self.assertEqual(1, session.number_of_requests_in_session)
+        self.assertEqual("some_product", product.name)
+
 
 if __name__ == "__main__":
     unittest.main()

pyravendb/tests/session_tests/test_store_entities.py

@@ -1,9 +1,16 @@
 from pyravendb.tests.test_base import TestBase
 from pyravendb.store.document_store import DocumentStore
 from pyravendb.custom_exceptions import exceptions
+from dataclasses import dataclass
 import unittest
 
 
+@dataclass()
+class Fish:
+    name: str
+    weight: int
+
+
 class Foo(object):
     def __init__(self, name, key):
         self.name = name
@@ -91,6 +98,18 @@ def test_store_the_same_documents_should_work(self):
             results = list(session.query().raw_query("From Foos"))
             self.assertEqual(len(results), 40 * 4)
 
+    def test_store_dataclass(self):
+        with self.store.open_session() as session:
+            fishie = Fish(name="Tuna", weight=100)
+            session.store(fishie, "fish/1")
+            session.save_changes()
+
+        with self.store.open_session() as session:
+            fishie = session.load("fish/1")
+            self.assertIsNotNone(fishie)
+            self.assertEqual(100, fishie.weight)
+            self.assertEqual("Tuna", fishie.name)
+
 
 if __name__ == "__main__":
     unittest.main()