GAE datastore/anom unique fields

unique.py

from anom import Model, Key, transactional, delete_multi

from models.exceptions import IntegrityError


class Unique(Model):
    """A model to store unique values.
    The only purpose of this model is to "reserve" values that must be unique
    within a given scope, as a workaround because datastore doesn't support
    the concept of uniqueness for entity properties.
    For example, suppose we have a model `User` with three properties that
    must be unique across a given group: `username`, `auth_id` and `email`::
        
        from anom import props, Model, Key
        
        class User(UniqueModelMixin, Model):
            unique_fields = ("username", "auth_id", "email")

            name = props.String(indexed=True)
            surname = props.String(optional=True)
            username = props.String(indexed=True)
            auth_id = props.String(indexed=True)
            email = props.String(indexed=True)
            first_name = props.String()

    To ensure property uniqueness when creating a new `User`, we first create
    `Unique` records for those properties, and if everything goes well we can
    save the new `User` record::
    
        User(**data).put()

    Based on the idea from http://squeeville.com/2009/01/30/add-a-unique-constraint-to-google-app-engine/
    """

    @classmethod
    def create(cls, value):
        """Creates a new unique value.
        :param value:
            The value to be unique, as a string.
            The value should include the scope in which the value must be
            unique (ancestor, namespace, kind and/or property name).
            For example, for a unique property `email` from kind `User`, the
            value can be `User.email:me@myself.com`. In this case `User.email`
            is the scope, and `me@myself.com` is the value to be unique.
        :returns:
            True if the unique value was created, False otherwise.
        """
        entity = cls(key=Key(cls, value))
        txn = lambda: entity.put() if not entity.key.get() else None
        return transactional()(txn)() is not None

    @classmethod
    def create_multi(cls, values):
        """Creates multiple unique values at once.
        :param values:
            A sequence of values to be unique. See :meth:`create`.
        :returns:
            A tuple (bool, list_of_keys). If all values were created, bool is
            True and list_of_keys is empty. If one or more values weren't
            created, bool is False and the list contains all the values that
            already existed in datastore during the creation attempt.
        """
        keys = [Key(cls, value) for value in values]

        # Create all records transactional.
        created = []
        entities = [cls(key=key) for key in keys]
        for entity in entities:
            func = lambda: entity.put() if not entity.key.get() else None
            key = transactional()(func)()
            if key:
                created.append(key.key)

        if created != keys:
            # A poor man's "rollback": delete all recently created records.
            delete_multi(created)
            return False, [k.id() for k in keys if k not in created]

        return True, []


class UniqueModelMixin:
    unique_fields = ()

    def pre_put_hook(self):
        """check that the email_id is not present already and raise error if so"""
        if not self.key.id_or_name:
            cls = type(self)
            cls_name = cls.__name__
            success, existing = Unique.create_multi([
                f"{cls_name}.{fld}:{getattr(self, fld)}" for fld in self.unique_fields
            ])
            for efld in existing:
                raise IntegrityError(f'Duplicate entry: {efld}')