appengine db extras

dbextras.py

from google.appengine.ext import db
from django.utils import simplejson


### Functions ###
def fetch_references(entities, attr):
    """Fetches and sets ReferenceProperty attributes in one db call"""
    keys = []
    for entity in entities:
        key = getattr(entity, '_' + attr)
        if key:
            keys.append(key)
    
    ref_keys = {}
    for ref_entity in db.get(set(keys)):
        key = ref_entity.key()
        ref_keys[key] = ref_entity
        
    for entity in entities:
        key = getattr(entity, '_' + attr)
        if key:
            setattr(entity, attr, ref_keys.get(key))
        
    return entities


### Properties ###
class KeyNameReferenceProperty(db.Property):
    def __init__(self, ref_class, **kwargs):
        self.ref_class = ref_class
        db.Property.__init__(self, **kwargs)
    def get_value_for_datastore(self, model_instance):
        model = db.Property.get_value_for_datastore(self, model_instance)
        return str(model.key().name())
    def make_value_from_datastore(self, value):
        return self.ref_class.get_by_key_name(value)
    data_type = basestring

class IdReferenceProperty(db.Property):
    def __init__(self, ref_class, **kwargs):
        self.ref_class = ref_class
        db.Property.__init__(self, **kwargs)
    def get_value_for_datastore(self, model_instance):
        model = db.Property.get_value_for_datastore(self, model_instance)
        return int(model.key().id())
    def make_value_from_datastore(self, value):
        return self.ref_class.get_by_id(value)
    data_type = int


class JSONProperty(db.Property):
    data_type = db.Text
    
    def get_value_for_datastore(self, model_instance):
        value = super(JSONProperty, self).get_value_for_datastore(model_instance)
        return db.Text(simplejson.dumps(value))
        
    def make_value_from_datastore(self, value):
        return simplejson.loads(value)