Syntactic sugar classes for ancestor queries

boot.py

@@ -71,6 +71,8 @@ def setup_env():
             else:
                 if name == 'webapp2':
                     extra_paths.append(root)
+                elif name == 'webob_0_9':
+                    extra_paths.append(root)
         sys.path = extra_paths + sys.path
         from google.appengine.api import apiproxy_stub_map
 

db/base.py

@@ -155,7 +155,8 @@ def _value_for_db(self, value, field, field_kind, db_type, lookup):
         if db_type == 'key':
 #            value = self._value_for_db_key(value, field_kind)
             try:
-                value = key_from_path(field.model._meta.db_table, value)
+                if not isinstance(value, Key):
+                    value = key_from_path(field.model._meta.db_table, value)
             except (BadArgumentError, BadValueError,):
                 raise DatabaseError("Only strings and positive integers "
                                     "may be used as keys on GAE.")

db/compiler.py

@@ -22,7 +22,7 @@
 
 from .db_settings import get_model_indexes
 from .expressions import ExpressionEvaluator
-from .utils import commit_locked
+from .utils import AncestorKey, commit_locked
 
 
 # Valid query types (a dictionary is used for speedy lookups).
@@ -86,6 +86,7 @@ def __init__(self, compiler, fields):
         self.included_pks = None
         self.excluded_pks = ()
         self.has_negated_exact_filter = False
+        self.ancestor_key = None
         self.ordering = []
         self.db_table = self.query.get_meta().db_table
         self.pks_only = (len(fields) == 1 and fields[0].primary_key)
@@ -190,6 +191,14 @@ def add_filter(self, field, lookup_type, negated, value):
         # Optimization: batch-get by key; this is only suitable for
         # primary keys, not for anything that uses the key type.
         if field.primary_key and lookup_type in ('exact', 'in'):
+            if lookup_type == 'exact' and isinstance(value, AncestorKey):
+                if negated:
+                    raise DatabaseError("You can't negate an ancestor operator.")
+                if self.ancestor_key is not None:
+                    raise DatabaseError("You can't use more than one ancestor operator.")
+                self.ancestor_key = value.key
+                return
+
             if self.included_pks is not None:
                 raise DatabaseError("You can't apply multiple AND "
                                     "filters on the primary key. "
@@ -311,6 +320,8 @@ def _make_entity(self, entity):
     def _build_query(self):
         for query in self.gae_query:
             query.Order(*self.ordering)
+            if self.ancestor_key:
+                query.Ancestor(self.ancestor_key)
         if len(self.gae_query) > 1:
             return MultiQuery(self.gae_query, self.ordering)
         return self.gae_query[0]
@@ -356,36 +367,41 @@ class SQLCompiler(NonrelCompiler):
 class SQLInsertCompiler(NonrelInsertCompiler, SQLCompiler):
 
     @safe_call
-    def insert(self, data, return_id=False):
+    def insert(self, data_list, return_id=False):
         opts = self.query.get_meta()
         unindexed_fields = get_model_indexes(self.query.model)['unindexed']
-        kwds = {'unindexed_properties': []}
-        properties = {}
-        for field, value in data.iteritems():
-
-            # The value will already be a db.Key, but the Entity
-            # constructor takes a name or id of the key, and will
-            # automatically create a new key if neither is given.
-            if field.primary_key:
-                if value is not None:
-                    kwds['id'] = value.id()
-                    kwds['name'] = value.name()
-
-            # GAE does not store empty lists (and even does not allow
-            # passing empty lists to Entity.update) so skip them.
-            elif isinstance(value, (tuple, list)) and not len(value):
-                continue
-
-            # Use column names as property names.
-            else:
-                properties[field.column] = value
+        unindexed_cols = [opts.get_field(name).column
+                          for name in unindexed_fields]
+
+        entity_list = []
+        for data in data_list:
+            properties = {}
+            kwds = {'unindexed_properties': unindexed_cols}
+            for column, value in data.items():
+                # The value will already be a db.Key, but the Entity
+                # constructor takes a name or id of the key, and will
+                # automatically create a new key if neither is given.
+                if column == opts.pk.column:
+                    if value is not None:
+                        kwds['id'] = value.id()
+                        kwds['name'] = value.name()
+                        kwds['parent'] = value.parent()
+
+                # GAE does not store empty lists (and even does not allow
+                # passing empty lists to Entity.update) so skip them.
+                elif isinstance(value, (tuple, list)) and not len(value):
+                    continue
+
+                # Use column names as property names.
+                else:
+                    properties[column] = value
 
-            if field in unindexed_fields:
-                kwds['unindexed_properties'].append(field.column)
+            entity = Entity(opts.db_table, **kwds)
+            entity.update(properties)
+            entity_list.append(entity)
 
-        entity = Entity(opts.db_table, **kwds)
-        entity.update(properties)
-        return Put(entity)
+        keys = Put(entity_list)
+        return keys[0] if isinstance(keys, list) else keys
 
 
 class SQLUpdateCompiler(NonrelUpdateCompiler, SQLCompiler):

db/creation.py

@@ -32,6 +32,13 @@ def db_type(self, field):
         field is to be indexed, and the "text" db_type (db.Text) if
         it's registered as unindexed.
         """
+        from djangoappengine.fields import DbKeyField
+
+        # DBKeyField reads/stores db.Key objects directly
+        # so its treated as a special case
+        if isinstance(field, DbKeyField):
+            return field.db_type(connection=self.connection)
+
         if self.connection.settings_dict.get('STORE_RELATIONS_AS_DB_KEYS'):
             if field.primary_key or field.rel is not None:
                 return 'key'
@@ -70,7 +77,8 @@ def _create_test_db(self, *args, **kw):
             stub_manager.activate_test_stubs()
 
     def _destroy_test_db(self, *args, **kw):
-        if self._had_test_stubs:
-            stub_manager.deactivate_test_stubs()
-            stub_manager.setup_stubs(self.connection)
-        del self._had_test_stubs
+        if hasattr(self, '_had_test_stubs'):
+            if self._had_test_stubs:
+                stub_manager.deactivate_test_stubs()
+                stub_manager.setup_stubs(self.connection)
+            del self._had_test_stubs

db/models/manager.py

@@ -0,0 +1,14 @@
+from django.db.models import Manager as _baseManager
+from djangoappengine.db.utils import as_ancestor
+from djangoappengine.db.models.query import QuerySet
+
+class Manager(_baseManager):
+
+    def get_query_set(self):
+        """Returns a new QuerySet object.  Subclasses can override this method
+        to easily customize the behavior of the Manager.
+        """
+        return QuerySet(self.model, using=self._db)
+
+    def ancestor(self, ancestor):
+        return self.get_query_set().ancestor(ancestor)

db/models/query.py

@@ -0,0 +1,17 @@
+from django.db.models.query import QuerySet as _baseQuerySet
+from djangoappengine.db.utils import as_ancestor
+
+class QuerySet(_baseQuerySet):
+    def ancestor(self, ancestor):
+        """
+        Returns a new QuerySet instance with the args ANDed to the existing
+        set.
+        """      
+        return self._filter_or_exclude(False, pk=as_ancestor(ancestor))
+
+class EmptyQuerySet(QuerySet):
+    def ancestor(self, *args, **kwargs):
+        """
+        Always returns EmptyQuerySet.
+        """
+        return self

db/utils.py

@@ -1,4 +1,6 @@
+from google.appengine.api.datastore import Key
 from google.appengine.datastore.datastore_query import Cursor
+
 from django.db import models, DEFAULT_DB_ALIAS
 
 try:
@@ -20,7 +22,7 @@ def get_cursor(queryset):
     # Evaluate QuerySet.
     len(queryset)
     cursor = getattr(queryset.query, '_gae_cursor', None)
-    return Cursor.to_websafe_string(cursor)
+    return Cursor.to_websafe_string(cursor) if cursor else None
 
 
 def set_cursor(queryset, start=None, end=None):
@@ -57,3 +59,35 @@ def _commit_locked(*args, **kw):
     if callable(func_or_using):
         return inner_commit_locked(func_or_using, DEFAULT_DB_ALIAS)
     return lambda func: inner_commit_locked(func, func_or_using)
+
+class AncestorKey(object):
+    def __init__(self, key):
+        self.key = key
+
+def as_ancestor(key_or_model):
+    if key_or_model is None:
+        raise ValueError("key_or_model must not be None")
+
+    if isinstance(key_or_model, models.Model):
+        key_or_model = Key.from_path(key_or_model._meta.db_table, key_or_model.pk)
+
+    return AncestorKey(key_or_model)
+
+def make_key(*args, **kwargs):
+    parent = kwargs.pop('parent', None)
+
+    if kwargs:
+        raise AssertionError('Excess keyword arguments; received %s' % kwargs)
+
+    if not args or len(args) % 2:
+        raise AssertionError('A non-zero even number of positional arguments is required; received %s' % args)
+
+    if isinstance(parent, models.Model):
+        parent = Key.from_path(parent._meta.db_table, parent.pk)
+
+    converted_args = []
+    for i in xrange(0, len(args), 2):
+        model, id_or_name = args[i:i+2]
+        converted_args.extend((model._meta.db_table, id_or_name))
+
+    return Key.from_path(*converted_args, parent=parent)

fields.py

@@ -0,0 +1,87 @@
+from django.core.exceptions import ValidationError
+from django.db import models
+from django.utils.encoding import smart_unicode
+
+from djangoappengine.db.utils import AncestorKey
+
+from google.appengine.api.datastore import Key, datastore_errors
+
+import logging
+
+class DbKeyField(models.Field):
+    description = "A field for native database key objects"
+    __metaclass__ = models.SubfieldBase
+
+    def __init__(self, *args, **kwargs):
+        kwargs['blank'] = True
+
+        self.parent_key_attname = kwargs.pop('parent_key_name', None)
+
+        if self.parent_key_attname is not None and kwargs.get('primary_key', None) is None:
+            raise ValueError("Primary key must be true to use parent_key_name")
+
+        super(DbKeyField, self).__init__(*args, **kwargs)
+
+    def contribute_to_class(self, cls, name):
+        if self.primary_key:
+            assert not cls._meta.has_auto_field, "A model can't have more than one auto field."
+            cls._meta.has_auto_field = True
+            cls._meta.auto_field = self
+
+            if self.parent_key_attname is not None:
+                def get_parent_key(instance, instance_type=None):
+                    if instance is None:
+                        return self
+
+                    return instance.__dict__.get(self.parent_key_attname)
+
+                def set_parent_key(instance, value):
+                    if instance is None:
+                        raise AttributeError("Attribute must be accessed via instance")
+
+                    if not isinstance(value, Key):
+                        raise ValueError("'%s' must be a Key" % self.parent_key_attname)
+
+                    instance.__dict__[self.parent_key_attname] = value
+
+                setattr(cls, self.parent_key_attname, property(get_parent_key, set_parent_key))
+
+        super(DbKeyField, self).contribute_to_class(cls, name)
+
+    def to_python(self, value):
+        if value is None:
+            return None
+        if isinstance(value, Key):
+            return value
+        if isinstance(value, basestring):
+            if len(value) == 0:
+                return None
+
+            try:
+                return Key(encoded=value)
+            except datastore_errors.BadKeyError:
+                return Key.from_path(self.model._meta.db_table, long(value))
+        if isinstance(value, (int, long)):
+            return Key.from_path(self.model._meta.db_table, value)
+
+        raise ValidationError("DbKeyField does not accept %s" % type(value))
+
+    def get_prep_value(self, value):
+        if isinstance(value, AncestorKey):
+            return value
+        return self.to_python(value)
+
+    def pre_save(self, model_instance, add):
+        value = super(DbKeyField, self).pre_save(model_instance, add)
+
+        if add and value is None and self.parent_key_attname is not None and hasattr(model_instance, self.parent_key_attname):
+            stashed_parent = getattr(model_instance, self.parent_key_attname)
+            value = Key.from_path(self.model._meta.db_table, 0, parent=stashed_parent)
+
+        return value
+
+    def formfield(self, **kwargs):
+        return None
+
+    def value_to_string(self, obj):
+        return smart_unicode(self._get_val_from_obj(obj))

test.py

@@ -0,0 +1,32 @@
+from django.test import TestCase
+
+from google.appengine.datastore import datastore_stub_util
+
+from db.stubs import stub_manager
+
+class GAETestCase(TestCase):
+    def _pre_setup(self):
+        """Performs any pre-test setup.
+            * Set the dev_appserver consistency state.
+        """
+        super(GAETestCase,self)._pre_setup()
+
+        if hasattr(self, 'consistency_probability'):
+            datastore = stub_manager.testbed.get_stub('datastore_v3')
+            self._orig_policy = datastore._consistency_policy
+
+            datastore.SetConsistencyPolicy(datastore_stub_util.PseudoRandomHRConsistencyPolicy(probability=self.consistency_probability))
+
+
+    def _post_teardown(self):
+        """ Performs any post-test things. This includes:
+
+            * Putting back the original ROOT_URLCONF if it was changed.
+            * Force closing the connection, so that the next test gets
+              a clean cursor.
+        """
+        if hasattr(self, '_orig_policy'):
+            datastore = stub_manager.testbed.get_stub('datastore_v3')
+            datastore.SetConsistencyPolicy(self._orig_policy)
+
+        super(GAETestCase,self)._post_teardown()

tests/__init__.py

@@ -1,8 +1,10 @@
 from .backend import BackendTest
+#from .decimals import DecimalTest
 from .field_db_conversion import FieldDBConversionTest
 from .field_options import FieldOptionsTest
 from .filter import FilterTest
-from .keys import KeysTest
+from .keys import KeysTest, DbKeyFieldTest, AncestorQueryTest, ParentKeyTest
 from .not_return_sets import NonReturnSetsTest
 from .order import OrderTest
 from .transactions import TransactionTest
+from .ancestor import AncestorTest