1# Copyright 2015 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Base classes for Model classes that can be internal-only."""
6
7from google.appengine.ext import ndb
8
9from dashboard import datastore_hooks
10
11
12class InternalOnlyModel(ndb.Model):
13  """A superclass for Models that have an internal_only property."""
14
15  @classmethod
16  def _post_get_hook(cls, key, future):  # pylint: disable=unused-argument
17    """Throws an exception when external users try to get() internal data."""
18    entity = future.get_result()
19    if entity is None:
20      return
21    # Internal-only objects should never be accessed by non-internal accounts!
22    if (getattr(entity, 'internal_only', False) and
23        not datastore_hooks.IsUnalteredQueryPermitted()):
24      # Keep info about the fact that we're doing an access check out of the
25      # callstack in case app engine shows it to the user.
26      assert False
27
28
29class CreateHookInternalOnlyModel(InternalOnlyModel):
30  """Base Model class which implements a create hook called CreateCallback."""
31
32  def __init__(self, *args, **kwargs):
33    super(CreateHookInternalOnlyModel, self).__init__(*args, **kwargs)
34    # This attribute is used to keep track of whether its the first time the
35    # entity has been created, so that the CreateCallback is only called once.
36    self._is_saved = False
37
38  def _post_put_hook(self, future):
39    """Invokes a callback upon creating a new entity."""
40    if not self._is_saved:
41      try:
42        future.check_success()
43      except Exception:  # pylint: disable=broad-except
44        # check_success throws an exception for failures, but the reference
45        # isn't explicit about what type of exception gets thrown.
46        pass
47      if callable(getattr(self, 'CreateCallback', None)):
48        self.CreateCallback()
49    self._is_saved = True
50