-
-
Save xuru/aa617d2de8fe3facea9d3a5d700b5d29 to your computer and use it in GitHub Desktop.
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
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.by_email
GSI toExternalUser
to ensure the model metaclass doesn't bind both? Use a differentname=
and set thehash_key='last_name'
pointing at the baseUserBase
. Might also be worth testing the explicit object reference withhash_key=UserBase.last_name
instead of the column'smodel_name
(although maybe that should fail...)engine.delete(*engine.scan(User))
to avoid traversing the whole list before deletingUser.created > created_after
). That should beUser.created < created_after
right? Looks like copy-paste from the next section.