Implement hooks.

src/google/cloud/ndb/key.py

@@ -86,6 +86,7 @@
 
 
 import base64
+import functools
 
 from google.cloud.datastore import _app_engine_key_pb2
 from google.cloud.datastore import key as _key_module
@@ -773,7 +774,6 @@ def get(
         """
         return self.get_async(_options=_options).result()
 
-    @tasklets.tasklet
     @_options.ReadOptions.options
     def get_async(
         self,
@@ -835,9 +835,23 @@ def get_async(
         """
         from google.cloud.ndb import model  # avoid circular import
 
-        entity_pb = yield _datastore_api.lookup(self._key, _options)
-        if entity_pb is not _datastore_api._NOT_FOUND:
-            return model._entity_from_protobuf(entity_pb)
+        cls = model.Model._kind_map.get(self.kind())
+
+        @tasklets.tasklet
+        def get():
+            if cls:
+                cls._pre_get_hook(self)
+
+            entity_pb = yield _datastore_api.lookup(self._key, _options)
+            if entity_pb is not _datastore_api._NOT_FOUND:
+                return model._entity_from_protobuf(entity_pb)
+
+        future = get()
+        if cls:
+            future.add_done_callback(
+                functools.partial(cls._post_get_hook, self)
+            )
+        return future
 
     @_options.Options.options
     def delete(
@@ -932,7 +946,20 @@ def delete_async(
                 set operations will be combined into a single set_multi
                 operation.
         """
-        return _datastore_api.delete(self._key, _options)
+        from google.cloud.ndb import model  # avoid circular import
+
+        cls = model.Model._kind_map.get(self.kind())
+        if cls:
+            cls._pre_delete_hook(self)
+
+        future = _datastore_api.delete(self._key, _options)
+
+        if cls:
+            future.add_done_callback(
+                functools.partial(cls._post_delete_hook, self)
+            )
+
+        return future
 
     @classmethod
     def from_old_key(cls, old_key):

src/google/cloud/ndb/model.py

@@ -4006,7 +4006,6 @@ def _put(
 
     put = _put
 
-    @tasklets.tasklet
     @_options.Options.options
     def _put_async(
         self,
@@ -4055,12 +4054,20 @@ def _put_async(
             tasklets.Future: The eventual result will be the key for the
                 entity. This is always a complete key.
         """
-        entity_pb = _entity_to_protobuf(self)
-        key_pb = yield _datastore_api.put(entity_pb, _options)
-        if key_pb:
-            ds_key = helpers.key_from_protobuf(key_pb)
-            self._key = key_module.Key._from_ds_key(ds_key)
-        return self._key
+
+        @tasklets.tasklet
+        def put(self):
+            self._pre_put_hook()
+            entity_pb = _entity_to_protobuf(self)
+            key_pb = yield _datastore_api.put(entity_pb, _options)
+            if key_pb:
+                ds_key = helpers.key_from_protobuf(key_pb)
+                self._key = key_module.Key._from_ds_key(ds_key)
+            return self._key
+
+        future = put(self)
+        future.add_done_callback(self._post_put_hook)
+        return future
 
     put_async = _put_async
 
@@ -4199,7 +4206,6 @@ def _allocate_ids(
     allocate_ids = _allocate_ids
 
     @classmethod
-    @tasklets.tasklet
     @_options.Options.options
     def _allocate_ids_async(
         cls,
@@ -4262,18 +4268,30 @@ def _allocate_ids_async(
         if not size:
             raise TypeError("Must pass non-zero 'size' to 'allocate_ids'")
 
-        kind = cls._get_kind()
-        keys = [
-            key_module.Key(kind, None, parent=parent)._key for _ in range(size)
-        ]
-        key_pbs = yield _datastore_api.allocate(keys, _options)
-        keys = tuple(
-            (
-                key_module.Key._from_ds_key(helpers.key_from_protobuf(key_pb))
-                for key_pb in key_pbs
+        @tasklets.tasklet
+        def allocate_ids():
+            cls._pre_allocate_ids_hook(size, max, parent)
+            kind = cls._get_kind()
+            keys = [
+                key_module.Key(kind, None, parent=parent)._key
+                for _ in range(size)
+            ]
+            key_pbs = yield _datastore_api.allocate(keys, _options)
+            keys = tuple(
+                (
+                    key_module.Key._from_ds_key(
+                        helpers.key_from_protobuf(key_pb)
+                    )
+                    for key_pb in key_pbs
+                )
             )
+            return keys
+
+        future = allocate_ids()
+        future.add_done_callback(
+            functools.partial(cls._post_allocate_ids_hook, size, max, parent)
         )
-        return keys
+        return future
 
     allocate_ids_async = _allocate_ids_async
 
@@ -4741,6 +4759,38 @@ def _to_dict(self, include=None, *, exclude=None):
 
     to_dict = _to_dict
 
+    @classmethod
+    def _pre_allocate_ids_hook(cls, size, max, parent):
+        pass
+
+    @classmethod
+    def _post_allocate_ids_hook(cls, size, max, parent, future):
+        pass
+
+    @classmethod
+    def _pre_delete_hook(self, key):
+        pass
+
+    @classmethod
+    def _post_delete_hook(self, key, future):
+        pass
+
+    @classmethod
+    def _pre_get_hook(self, key):
+        pass
+
+    @classmethod
+    def _post_get_hook(self, key, future):
+        pass
+
+    @classmethod
+    def _pre_put_hook(self):
+        pass
+
+    @classmethod
+    def _post_put_hook(self, future):
+        pass
+
 
 class Expando(Model):
     __slots__ = ()

tests/unit/test_key.py

@@ -536,19 +536,56 @@ def test_urlsafe():
     @unittest.mock.patch("google.cloud.ndb.key._datastore_api")
     @unittest.mock.patch("google.cloud.ndb.model._entity_from_protobuf")
     def test_get(_entity_from_protobuf, _datastore_api):
+        class Simple(model.Model):
+            pass
+
         ds_future = tasklets.Future()
         ds_future.set_result("ds_entity")
         _datastore_api.lookup.return_value = ds_future
         _entity_from_protobuf.return_value = "the entity"
 
-        key = key_module.Key("a", "b", app="c")
+        key = key_module.Key("Simple", "b", app="c")
+        assert key.get() == "the entity"
+
+        _datastore_api.lookup.assert_called_once_with(
+            key._key, _options.ReadOptions()
+        )
+        _entity_from_protobuf.assert_called_once_with("ds_entity")
+
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @unittest.mock.patch("google.cloud.ndb.key._datastore_api")
+    @unittest.mock.patch("google.cloud.ndb.model._entity_from_protobuf")
+    def test_get_w_hooks(_entity_from_protobuf, _datastore_api):
+        class Simple(model.Model):
+            pre_get_calls = []
+            post_get_calls = []
+
+            @classmethod
+            def _pre_get_hook(cls, *args, **kwargs):
+                cls.pre_get_calls.append((args, kwargs))
+
+            @classmethod
+            def _post_get_hook(cls, key, future, *args, **kwargs):
+                assert isinstance(future, tasklets.Future)
+                cls.post_get_calls.append(((key,) + args, kwargs))
+
+        ds_future = tasklets.Future()
+        ds_future.set_result("ds_entity")
+        _datastore_api.lookup.return_value = ds_future
+        _entity_from_protobuf.return_value = "the entity"
+
+        key = key_module.Key("Simple", 42)
         assert key.get() == "the entity"
 
         _datastore_api.lookup.assert_called_once_with(
             key._key, _options.ReadOptions()
         )
         _entity_from_protobuf.assert_called_once_with("ds_entity")
 
+        assert Simple.pre_get_calls == [((key,), {})]
+        assert Simple.post_get_calls == [((key,), {})]
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     @unittest.mock.patch("google.cloud.ndb.key._datastore_api")
@@ -584,16 +621,49 @@ def test_get_async_not_found(_datastore_api):
     @pytest.mark.usefixtures("in_context")
     @unittest.mock.patch("google.cloud.ndb.key._datastore_api")
     def test_delete(_datastore_api):
+        class Simple(model.Model):
+            pass
+
         future = tasklets.Future()
         _datastore_api.delete.return_value = future
         future.set_result("result")
 
-        key = key_module.Key("a", "b", app="c")
+        key = key_module.Key("Simple", "b", app="c")
         assert key.delete() == "result"
         _datastore_api.delete.assert_called_once_with(
             key._key, _options.Options()
         )
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @unittest.mock.patch("google.cloud.ndb.key._datastore_api")
+    def test_delete_w_hooks(_datastore_api):
+        class Simple(model.Model):
+            pre_delete_calls = []
+            post_delete_calls = []
+
+            @classmethod
+            def _pre_delete_hook(cls, *args, **kwargs):
+                cls.pre_delete_calls.append((args, kwargs))
+
+            @classmethod
+            def _post_delete_hook(cls, key, future, *args, **kwargs):
+                assert isinstance(future, tasklets.Future)
+                cls.post_delete_calls.append(((key,) + args, kwargs))
+
+        future = tasklets.Future()
+        _datastore_api.delete.return_value = future
+        future.set_result("result")
+
+        key = key_module.Key("Simple", 42)
+        assert key.delete() == "result"
+        _datastore_api.delete.assert_called_once_with(
+            key._key, _options.Options()
+        )
+
+        assert Simple.pre_delete_calls == [((key,), {})]
+        assert Simple.post_delete_calls == [((key,), {})]
+
     @staticmethod
     @unittest.mock.patch("google.cloud.ndb.key._datastore_api")
     def test_delete_in_transaction(_datastore_api, in_context):

tests/unit/test_model.py

@@ -2913,6 +2913,36 @@ def test__put_async(_datastore_api):
             entity_pb, _options.Options()
         )
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @unittest.mock.patch("google.cloud.ndb.model._datastore_api")
+    def test__put_w_hooks(_datastore_api):
+        class Simple(model.Model):
+            def __init__(self):
+                super(Simple, self).__init__()
+                self.pre_put_calls = []
+                self.post_put_calls = []
+
+            def _pre_put_hook(self, *args, **kwargs):
+                self.pre_put_calls.append((args, kwargs))
+
+            def _post_put_hook(self, future, *args, **kwargs):
+                assert isinstance(future, tasklets.Future)
+                self.post_put_calls.append((args, kwargs))
+
+        entity = Simple()
+        _datastore_api.put.return_value = future = tasklets.Future()
+        future.set_result(None)
+
+        entity_pb = model._entity_to_protobuf(entity)
+        assert entity._put() == entity.key
+        _datastore_api.put.assert_called_once_with(
+            entity_pb, _options.Options()
+        )
+
+        assert entity.pre_put_calls == [((), {})]
+        assert entity.post_put_calls == [((), {})]
+
     @staticmethod
     def test__lookup_model():
         class ThisKind(model.Model):
@@ -3043,6 +3073,56 @@ class Simple(model.Model):
         ]
         assert call_options == _options.Options()
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @unittest.mock.patch("google.cloud.ndb.model._datastore_api")
+    def test_allocate_ids_w_hooks(_datastore_api):
+        completed = [
+            entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="Simple", id=21)],
+            ),
+            entity_pb2.Key(
+                partition_id=entity_pb2.PartitionId(project_id="testing"),
+                path=[entity_pb2.Key.PathElement(kind="Simple", id=42)],
+            ),
+        ]
+        _datastore_api.allocate.return_value = utils.future_result(completed)
+
+        class Simple(model.Model):
+            pre_allocate_id_calls = []
+            post_allocate_id_calls = []
+
+            @classmethod
+            def _pre_allocate_ids_hook(cls, *args, **kwargs):
+                cls.pre_allocate_id_calls.append((args, kwargs))
+
+            @classmethod
+            def _post_allocate_ids_hook(
+                cls, size, max, parent, future, *args, **kwargs
+            ):
+                assert isinstance(future, tasklets.Future)
+                cls.post_allocate_id_calls.append(
+                    ((size, max, parent) + args, kwargs)
+                )
+
+        keys = Simple.allocate_ids(2)
+        assert keys == (
+            key_module.Key("Simple", 21),
+            key_module.Key("Simple", 42),
+        )
+
+        call_keys, call_options = _datastore_api.allocate.call_args[0]
+        call_keys = [key_module.Key._from_ds_key(key) for key in call_keys]
+        assert call_keys == [
+            key_module.Key("Simple", None),
+            key_module.Key("Simple", None),
+        ]
+        assert call_options == _options.Options()
+
+        assert Simple.pre_allocate_id_calls == [((2, None, None), {})]
+        assert Simple.post_allocate_id_calls == [((2, None, None), {})]
+
     @staticmethod
     @pytest.mark.usefixtures("in_context")
     def test_allocate_ids_with_max():