Testing some inheritance flows

bloop_test.py

import random
import uuid
from string import ascii_letters
from enum import Enum

import boto3
from datetime import datetime, timezone
from bloop import (
    BaseModel, Column, String, UUID, GlobalSecondaryIndex, LocalSecondaryIndex, Engine, DateTime, Boolean,
    Set
)


class StringEnum(String):
    def __init__(self, enum_cls):
        self.enum_cls = enum_cls
        super().__init__()

    def dynamo_dump(self, value, *, context, **kwargs):
        if value is None:
            return value
        value = value.name
        return super().dynamo_dump(value, context=context, **kwargs)

    def dynamo_load(self, value, *, context, **kwargs):
        if value is None:
            return value
        value = super().dynamo_load(value, context=context, **kwargs)
        return self.enum_cls[value]


class Role(Enum):
    user = "user"
    curator = "curator"
    superuser = "super_user"
    admin = "admin"


# ------------------------------------------------------------------------------
# Mixins
# ------------------------------------------------------------------------------
class IdentityMixin(object):
    roles = Column(Set(StringEnum(Role)))

    @property
    def is_active(self):
        return True

    @property
    def is_authenticated(self):
        return True

    @property
    def is_anonymous(self):
        return False

    def get_id(self):
        return self.id


class UUIDHashKey(object):
    id = Column(UUID, hash_key=True)


class CreatedRangeKey(object):
    created = Column(DateTime, range_key=True)
    by_created = LocalSecondaryIndex(projection="all", range_key=created)


# ------------------------------------------------------------------------------
# Models
# NOTE: Notice that the above mixins do not have an abstract flag.  It's not needed
# because we don't bind them (we bind all classes that are derived from BaseModel).
# Also, notice that I only set the abstract flag on classes that I don't want a table
# for (MyBaseModel), and I derive ExternalUser from User now.
# ------------------------------------------------------------------------------
class MyBaseModel(BaseModel, UUIDHashKey):
    class Meta:
        abstract = True
    updated = Column(DateTime)
    active = Column(Boolean)


class User(MyBaseModel, CreatedRangeKey, IdentityMixin):
    first_name = Column(String)
    last_name = Column(String)
    email = Column(String)
    by_email = GlobalSecondaryIndex(projection='all', name='email-index', hash_key='last_name')

    def __str__(self):
        return "{} {}: {}".format(self.first_name, self.last_name, self.email)


class ExternalUser(User):
    company = Column(String)
    by_email = GlobalSecondaryIndex(projection='all', name='email-index-two', hash_key=User.email)


print("ExternalUser.by_email: {}".format(ExternalUser.by_email))

# ------------------------------------------------------------------------------
# Bind
# ------------------------------------------------------------------------------
dynamodb_local = boto3.client("dynamodb", endpoint_url="http://127.0.0.1:8000")
engine = Engine(dynamodb=dynamodb_local)
engine.bind(BaseModel)

# ------------------------------------------------------------------------------
# create some instances
# ------------------------------------------------------------------------------


def create_objs(cls, **extra):
    now = datetime.now(timezone.utc)

    first = "".join([random.choice(ascii_letters) for x in range(8)])
    last = "".join([random.choice(ascii_letters) for x in range(12)])
    email = first + "." + last + "@example.com"

    obj = cls(
        id=uuid.uuid4(),
        created=now,
        updated=now,
        first_name=first,
        last_name=last,
        email=email,
        **extra
    )
    engine.save(obj)
    return obj

before_saved_id = None
for x in range(10):
    user = create_objs(ExternalUser, **{'company': 'Acme', 'roles': [Role.user, Role.admin]})
    before_saved_id = user.id


for x in range(10):
    user = create_objs(User, **{'roles': [Role.user]})

after_saved_id = None
created_after = datetime.now(timezone.utc)
for x in range(10):
    user = create_objs(ExternalUser, **{'company': 'Acme', 'roles': [Role.user, Role.admin]})
    after_saved_id = user.id

# ------------------------------------------------------------------------------
# Scan users and external users
# ------------------------------------------------------------------------------
print("Users:")
for user in engine.scan(User):
    print("  {}".format(user))

print("External Users:")
for user in engine.scan(ExternalUser):
    print("  [{}] {}".format(user.company, user))

# ------------------------------------------------------------------------------
# User GSI and LSI
# ------------------------------------------------------------------------------
# get the first external user
user = engine.scan(ExternalUser).first()

# use the users email to find him using the GSI
print("Get external user using GSI")
cond = ExternalUser.email == user.email
user = engine.query(ExternalUser.by_email, key=cond).one()
print("Found: {}".format(user))


# Try to find the 'before_saved_id' user with a created date that is after it was created.  Should return no users.
cond = (User.id == before_saved_id) & (User.created > created_after)
q = engine.query(User.by_created, key=cond)
print("External user created during first batch:")
for user in q:
    print(user)

# Try to find the 'after_saved_id' user with a created date that is before it was created.  Should return 1 user.
cond = (User.id == after_saved_id) & (User.created > created_after)
q = engine.query(User.by_created, key=cond)
print("External users created after first batch:")
for user in q:
    print(user)


# ------------------------------------------------------------------------------
# Cleanup
# ------------------------------------------------------------------------------
for user in engine.scan(User):
    engine.delete(user)
for user in engine.scan(ExternalUser):
    engine.delete(user)

The LSI CreatedRangeKey.by_created has the same range key as the table, so it should be equivalent to a query on the table. Does DynamoDB actually let you create this? If not, bloop should validate against it and fail fast.

Yeah, you're right. For some reason, it seems to work. May have to look into that.

Can you add a different by_email GSI to ExternalUser to ensure the model metaclass doesn't bind both? Use a different name= and set the hash_key='last_name' pointing at the base UserBase. Might also be worth testing the explicit object reference with hash_key=UserBase.last_name instead of the column's model_name (although maybe that should fail...)

Yeah, this seems to work (will bind the last one), but interestingly, putting a base class's attribute in for the hask_key works as well.

nit: You can simplify to engine.delete(*engine.scan(User)) to avoid traversing the whole list before deleting

Is there a way to batch delete? Passing in a list doesn't work.

The comment on 174 says "created before the created_after date" but then the condition on 175 says User.created > created_after). That should be User.created < created_after right? Looks like copy-paste from the next section.

Yeah, will fix