def test_serialize_deserialize(self, mock_time):
mock_time.side_effect = [1559692800, 1559692800] # 2019-06-05 00:00:00 UTC
model = AttributeTestModel()
model.ttl_attr = timedelta(minutes=1)
assert model.ttl_attr == datetime(2019, 6, 5, 0, 1, tzinfo=timezone.utc)
s = TTLAttribute().serialize(model.ttl_attr)
assert s == '1559692860'
assert TTLAttribute().deserialize(s) == datetime(2019, 6, 5, 0, 1, 0, tzinfo=timezone.utc)
class Base(Model):
class Meta:
table_name = "appstore"
region = "us-west-2"
billing_mode = "PAY_PER_REQUEST"
pk = CompoundTemplateAttribute(
hash_key=True,
template="$type_#$ulid",
attrs=["type_", "ulid"],
)
ulid = ULIDAttribute()
cls = DiscriminatorAttribute(attr_name="__type")
ttl = TTLAttribute(null=True)
gsi1 = GSI1()
gsi2 = GSI2()
@property
def type_(self):
return self.__class__.__name__.upper()
@property
def created_at(self) -> datetime:
"""Based on a ULID `pk`, extract an aware datetime"""
if ts := getattr(self.pk, "timestamp", None):
return ts().datetime
if ts := getattr(self.ulid, "timestamp", None):
return ts().datetime
def test_default_and_default_for_new(self):
with pytest.raises(
ValueError,
match=
'An attribute cannot have both default and default_for_new parameters'
):
TTLAttribute(default=timedelta(seconds=1),
default_for_new=timedelta(seconds=2))
class AttributeTestModel(Model):
class Meta:
host = 'http://localhost:8000'
table_name = 'test'
binary_attr = BinaryAttribute(hash_key=True)
binary_set_attr = BinarySetAttribute()
number_attr = NumberAttribute()
number_set_attr = NumberSetAttribute()
unicode_attr = UnicodeAttribute()
unicode_set_attr = UnicodeSetAttribute()
datetime_attr = UTCDateTimeAttribute()
bool_attr = BooleanAttribute()
json_attr = JSONAttribute()
map_attr = MapAttribute()
ttl_attr = TTLAttribute()
class EventLog(Model):
"""Event processing log"""
class Meta:
table_name = os.environ.get("EVENT_LOG_TABLE", "event-log")
host = os.environ.get('DYNAMODB_HOST')
region = Session().get_config_variable("region")
s3_key = UnicodeAttribute(hash_key=True)
s3_bucket = UnicodeAttribute(attr_name='bucket')
status = UnicodeAttribute()
gzip = BooleanAttribute(default=False)
function = UnicodeAttribute()
region = UnicodeAttribute()
functional_key_name = UnicodeAttribute(null=True, attr_name="funcKeyName")
functional_key_value = UnicodeAttribute(null=True, attr_name="funcKeyVal")
trace_id = UnicodeAttribute(null=True, attr_name="trace")
received_time = UTCDateTimeAttribute(attr_name="received")
processed_time = UTCDateTimeAttribute(null=True, attr_name="processed")
error = UnicodeAttribute(null=True)
expires = TTLAttribute(default_for_new=EVENT_LOG_TTL)
status_index = StatusIndex()
def set_functional_key(self, functional_key_name: str,
functional_key_value: str):
self.functional_key_name = functional_key_name
self.functional_key_value = functional_key_value
def mark_processed(self):
if not self.processed_time:
self.processed_time = datetime.now(timezone.utc)
self.status = STATUS_DONE
self.save()
if getattr(self, 'logger'):
self.logger.info('processed s3 object')
def mark_failed(self, error: str):
if not self.processed_time:
self.processed_time = datetime.now(timezone.utc)
self.error = str(error)
self.status = STATUS_FAILED
self.save()
if getattr(self, 'logger'):
self.logger.warning(f'processing s3 object failed: {error}')
class Server(Model):
class OnlineIndex(GlobalSecondaryIndex):
class Meta:
read_capacity_units = 3
write_capacity_units = 3
projection = AllProjection()
online = BooleanAsNumberAttribute(hash_key=True)
class TimeLastSeenIndex(GlobalSecondaryIndex):
class Meta:
read_capacity_units = 3
write_capacity_units = 3
projection = AllProjection()
online = BooleanAsNumberAttribute(hash_key=True)
time_last_seen = NumberAttribute(range_key=True)
class Meta:
table_name = "MSU-server"
read_capacity_units = 3
write_capacity_units = 3
session_key = NumberAttribute(hash_key=True)
token = NumberAttribute()
online = BooleanAsNumberAttribute(default=False)
online_view = OnlineIndex()
ipv4 = ServerIpMap(null=True)
ipv6 = ServerIpMap(null=True)
info = InfoMap(null=True)
time_first_seen = NumberAttribute(null=True)
time_last_seen = NumberAttribute(null=True)
time_last_seen_view = TimeLastSeenIndex()
ttl = TTLAttribute()
def __repr__(self):
return str(vars(self))
class Url(Model):
"""
DynamoDB Model for url storage
"""
class Meta:
table_name = 'MagicUrl-urls'
read_capacity_units = 1
write_capacity_units = 1
class CodeIndex(GlobalSecondaryIndex):
class Meta:
read_capacity_units = 1
write_capacity_units = 1
projection = AllProjection()
code = UnicodeAttribute(hash_key=True)
url = UnicodeAttribute(hash_key=True)
code = UnicodeAttribute()
code_index = CodeIndex()
ttl = TTLAttribute()
def save(self, **kwargs):
"""
Generates the shortened code before saving
"""
# Set automatic code
self.code = b64encode(
md5(self.url.encode('utf-8')).hexdigest()[-4:].encode(
'utf-8')).decode('utf-8').replace('=', '').replace('/', '_')
# Set expire time
today = datetime.utcnow().replace(tzinfo=pytz.utc)
self.ttl = today + timedelta(days=7)
# Call parent
super(Url, self).save(**kwargs)
class IpPort(Model):
class OnlineIndex(GlobalSecondaryIndex):
class Meta:
read_capacity_units = 3
write_capacity_units = 3
projection = AllProjection()
online = BooleanAsNumberAttribute(hash_key=True)
class TimeLastSeenIndex(GlobalSecondaryIndex):
class Meta:
read_capacity_units = 3
write_capacity_units = 3
projection = AllProjection()
online = BooleanAsNumberAttribute(hash_key=True)
time_last_seen = NumberAttribute(range_key=True)
class Meta:
table_name = "MSU-ip-port"
read_capacity_units = 3
write_capacity_units = 3
server_id = UnicodeAttribute(hash_key=True)
session_key = NumberAttribute()
server_ip = ServerIpMap()
online = BooleanAsNumberAttribute(default=False)
online_view = OnlineIndex()
time_last_seen = NumberAttribute(null=True)
time_last_seen_view = TimeLastSeenIndex()
ttl = TTLAttribute()
def __repr__(self):
return str(vars(self))
def test_serialize_none(self):
model = AttributeTestModel()
model.ttl_attr = None
assert model.ttl_attr == None
assert TTLAttribute().serialize(model.ttl_attr) == None
def test_serialize_timedelta(self, mock_time):
mock_time.side_effect = [1559692800] # 2019-06-05 00:00:00 UTC
assert TTLAttribute().serialize(
timedelta(seconds=60)) == str(1559692800 + 60)
class OAuthToken(Model):
"""
Stores token data from some OAuth provider
"""
ALIVE = ALIVE
DEAD = DEAD
access_token = UnicodeAttribute(hash_key=True)
refresh_token = UnicodeAttribute(null=True)
token_type = UnicodeAttribute(null=True)
# Consistency would be too easy
expires_in = NumberAttribute(null=True)
expires_at = NumberAttribute(null=True)
scope = UnicodeAttribute(null=True)
user_info = JSONAttribute(null=True)
state_index = TokensByState()
state = UnicodeAttribute(default=ALIVE)
updated_at = UTCDateTimeAttribute()
created_at = UTCDateTimeAttribute(range_key=True)
ttl = TTLAttribute(null=True)
def save(self, *args, **kwargs):
timestamp = timezone.now()
if not self.created_at:
self.created_at = timestamp
self.updated_at = timestamp
super().save(*args, **kwargs)
def update(self, *args, **kwargs):
self.updated_at = timezone.now()
super().save(*args, **kwargs)
def set_updated_at(self, *args, **kwargs):
self.updated_at = timezone.now()
def set_expiration(self):
self.ttl = datetime.timedelta(days=30)
@property
def session_data(self):
return dict(
access_token=self.access_token,
token_type=self.token_type,
refresh_token=self.refresh_token,
expires_at=self.expires_at,
expires_in=self.expires_in,
)
@classmethod
def create_if_non_existent(cls):
if not cls.exists():
cls.create_table(read_capacity_units=1, write_capacity_units=1, wait=True)
class Meta:
table_name = settings.OAUTH_TOKEN_TABLE_NAME
region = settings.AWS_REGION
