Implement cache policy. (googleapis#116)

docs/conf.py

@@ -41,6 +41,7 @@
     ("py:class", "google.cloud.datastore_v1.proto.entity_pb2.Entity"),
     ("py:class", "_datastore_query.Cursor"),
     ("py:meth", "_datastore_query.Cursor.urlsafe"),
+    ("py:class", "google.cloud.ndb.context._Context"),
     ("py:class", "google.cloud.ndb.metadata._BaseMetadata"),
     ("py:class", "google.cloud.ndb._options.ReadOptions"),
     ("py:class", "QueryIterator"),

docs/context.rst

@@ -0,0 +1,7 @@
+#######
+Context
+#######
+
+.. automodule:: google.cloud.ndb.context
+    :members:
+    :show-inheritance:

docs/index.rst

@@ -7,6 +7,7 @@
    :maxdepth: 2
 
    client
+   context
    key
    model
    query

src/google/cloud/ndb/client.py

@@ -92,7 +92,7 @@ def __init__(self, project=None, namespace=None, credentials=None):
         self.secure = not emulator
 
     @contextlib.contextmanager
-    def context(self):
+    def context(self, cache_policy=None):
         """Establish a context for a set of NDB calls.
 
         This method provides a context manager which establishes the runtime
@@ -121,8 +121,13 @@ def context(self):
         In a web application, it is recommended that a single context be used
         per HTTP request. This can typically be accomplished in a middleware
         layer.
+
+        Arguments:
+            cache_policy (Optional[Callable[[key.Key], bool]]): The
+                cache policy to use in this context. See:
+                :meth:`~google.cloud.ndb.context.Context.set_cache_policy`.
         """
-        context = context_module.Context(self)
+        context = context_module.Context(self, cache_policy=cache_policy)
         with context.use():
             yield context
 

src/google/cloud/ndb/context.py

@@ -21,6 +21,7 @@
 from google.cloud.ndb import _datastore_api
 from google.cloud.ndb import _eventloop
 from google.cloud.ndb import exceptions
+from google.cloud.ndb import model
 
 
 __all__ = [
@@ -32,20 +33,6 @@
 ]
 
 
-_ContextTuple = collections.namedtuple(
-    "_ContextTuple",
-    [
-        "client",
-        "eventloop",
-        "stub",
-        "batches",
-        "commit_batches",
-        "transaction",
-        "cache",
-    ],
-)
-
-
 class _LocalState(threading.local):
     """Thread local state."""
 
@@ -97,6 +84,41 @@ def get_and_validate(self, key):
             raise KeyError(key)
 
 
+def _default_cache_policy(key):
+    """The default cache policy.
+
+    Defers to ``_use_cache`` on the Model class for the key's kind.
+
+    See: :meth:`~google.cloud.ndb.context.Context.set_cache_policy`
+    """
+    flag = None
+    if key is not None:
+        modelclass = model.Model._kind_map.get(key.kind())
+        if modelclass is not None:
+            policy = getattr(modelclass, "_use_cache", None)
+            if policy is not None:
+                if isinstance(policy, bool):
+                    flag = policy
+                else:
+                    flag = policy(key)
+
+    return flag
+
+
+_ContextTuple = collections.namedtuple(
+    "_ContextTuple",
+    [
+        "client",
+        "eventloop",
+        "stub",
+        "batches",
+        "commit_batches",
+        "transaction",
+        "cache",
+    ],
+)
+
+
 class _Context(_ContextTuple):
     """Current runtime state.
 
@@ -106,8 +128,8 @@ class _Context(_ContextTuple):
     loop. A new context can be derived from an existing context using
     :meth:`new`.
 
-    :class:`Context` is a subclass of :class:`_Context` which provides
-    only publicly facing interface. The use of two classes is only to provide a
+    :class:`Context` is a subclass of :class:`_Context` which provides only
+    publicly facing interface. The use of two classes is only to provide a
     distinction between public and private API.
 
     Arguments:
@@ -123,6 +145,7 @@ def __new__(
         commit_batches=None,
         transaction=None,
         cache=None,
+        cache_policy=None,
     ):
         if eventloop is None:
             eventloop = _eventloop.EventLoop()
@@ -145,7 +168,7 @@ def __new__(
         else:
             cache = _Cache()
 
-        return super(_Context, cls).__new__(
+        context = super(_Context, cls).__new__(
             cls,
             client=client,
             eventloop=eventloop,
@@ -156,6 +179,10 @@ def __new__(
             cache=cache,
         )
 
+        context.set_cache_policy(cache_policy)
+
+        return context
+
     def new(self, **kwargs):
         """Create a new :class:`_Context` instance.
 
@@ -205,9 +232,9 @@ def get_cache_policy(self):
             Callable: A function that accepts a
                 :class:`~google.cloud.ndb.key.Key` instance as a single
                 positional argument and returns a ``bool`` indicating if it
-                should be cached.  May be :data:`None`.
+                should be cached. May be :data:`None`.
         """
-        raise NotImplementedError
+        return self.cache_policy
 
     def get_datastore_policy(self):
         """Return the current context datastore policy function.
@@ -238,7 +265,7 @@ def get_memcache_timeout_policy(self):
             Callable: A function that accepts a
                 :class:`~google.cloud.ndb.key.Key` instance as a single
                 positional argument and returns an ``int`` indicating the
-                timeout, in seconds, for the key. :data:`0` implies the default
+                timeout, in seconds, for the key. ``0`` implies the default
                 timeout. May be :data:`None`.
         """
         raise NotImplementedError
@@ -252,7 +279,16 @@ def set_cache_policy(self, policy):
                 positional argument and returns a ``bool`` indicating if it
                 should be cached.  May be :data:`None`.
         """
-        raise NotImplementedError
+        if policy is None:
+            policy = _default_cache_policy
+
+        elif isinstance(policy, bool):
+            flag = policy
+
+            def policy(key):
+                return flag
+
+        self.cache_policy = policy
 
     def set_datastore_policy(self, policy):
         """Set the context datastore policy function.
@@ -283,7 +319,7 @@ def set_memcache_timeout_policy(self, policy):
             policy (Callable): A function that accepts a
                 :class:`~google.cloud.ndb.key.Key` instance as a single
                 positional argument and returns an ``int`` indicating the
-                timeout, in seconds, for the key. :data:`0` implies the default
+                timeout, in seconds, for the key. ``0`` implies the default
                 timout. May be :data:`None`.
         """
         raise NotImplementedError
@@ -319,59 +355,45 @@ def in_transaction(self):
         """
         return self.transaction is not None
 
-    @staticmethod
-    def default_cache_policy(key):
-        """Default cache policy.
-
-        This defers to :meth:`~google.cloud.ndb.model.Model._use_cache`.
-
-        Args:
-            key (google.cloud.ndb.model.key.Key): The key.
-
-        Returns:
-            Union[bool, NoneType]: Whether to cache the key.
-        """
-        raise NotImplementedError
-
     @staticmethod
     def default_datastore_policy(key):
         """Default cache policy.
 
-        This defers to :meth:`~google.cloud.ndb.model.Model._use_datastore`.
+        This defers to ``Model._use_datastore``.
 
         Args:
-            key (google.cloud.ndb.model.key.Key): The key.
+            key (google.cloud.ndb.key.Key): The key.
 
         Returns:
-            Union[bool, NoneType]: Whether to use datastore.
+            Union[bool, None]: Whether to use datastore.
         """
         raise NotImplementedError
 
     @staticmethod
     def default_memcache_policy(key):
         """Default memcache policy.
 
-        This defers to :meth:`~google.cloud.ndb.model.Model._use_memcache`.
+        This defers to ``Model._use_memcache``.
 
         Args:
-            key (google.cloud.ndb.model.key.Key): The key.
+            key (google.cloud.ndb.key.Key): The key.
 
         Returns:
-            Union[bool, NoneType]: Whether to cache the key.
+            Union[bool, None]: Whether to cache the key.
         """
         raise NotImplementedError
 
     @staticmethod
     def default_memcache_timeout_policy(key):
         """Default memcache timeout policy.
 
-        This defers to :meth:`~google.cloud.ndb.model.Model._memcache_timeout`.
+        This defers to ``Model._memcache_timeout``.
 
         Args:
-            key (google.cloud.ndb.model.key.Key): The key.
+            key (google.cloud.ndb.key.Key): The key.
 
         Returns:
-            Union[int, NoneType]: Memcache timeout to use.
+            Union[int, None]: Memcache timeout to use.
         """
         raise NotImplementedError
 
@@ -416,6 +438,15 @@ def urlfetch(self, *args, **kwargs):
         """Fetch a resource using HTTP."""
         raise NotImplementedError
 
+    def _use_cache(self, key, options):
+        """Return whether to use the context cache for this key."""
+        flag = options.use_cache
+        if flag is None:
+            flag = self.cache_policy(key)
+        if flag is None:
+            flag = True
+        return flag
+
 
 class ContextOptions:
     __slots__ = ()

src/google/cloud/ndb/key.py

@@ -842,12 +842,13 @@ def get_async(
 
         @tasklets.tasklet
         def get():
-            if _options.use_cache:
+            context = context_module.get_context()
+            use_cache = context._use_cache(self, _options)
+
+            if use_cache:
                 try:
                     # This result may be None, if None is cached for this key.
-                    return context_module.get_context().cache.get_and_validate(
-                        self
-                    )
+                    return context.cache.get_and_validate(self)
                 except KeyError:
                     pass
 
@@ -857,8 +858,8 @@ def get():
             else:
                 result = None
 
-            if _options.use_cache:
-                context_module.get_context().cache[self] = result
+            if use_cache:
+                context.cache[self] = result
 
             return result
 
@@ -971,8 +972,11 @@ def delete_async(
         @tasklets.tasklet
         def delete():
             result = yield _datastore_api.delete(self._key, _options)
-            if _options.use_cache:
-                context_module.get_context().cache[self] = None
+
+            context = context_module.get_context()
+            if context._use_cache(self, _options):
+                context.cache[self] = None
+
             return result
 
         future = delete()

src/google/cloud/ndb/model.py

@@ -4751,8 +4751,9 @@ def put(self):
                 ds_key = helpers.key_from_protobuf(key_pb)
                 self._key = key_module.Key._from_ds_key(ds_key)
 
-                if _options.use_cache:
-                    context_module.get_context().cache[self._key] = self
+                context = context_module.get_context()
+                if context._use_cache(self._key, _options):
+                    context.cache[self._key] = self
 
             return self._key
 

tests/system/conftest.py

@@ -92,5 +92,5 @@ def namespace():
 @pytest.fixture
 def client_context(namespace):
     client = ndb.Client(namespace=namespace)
-    with client.context():
-        yield
+    with client.context(cache_policy=False) as the_context:
+        yield the_context