Skip to content

Commit

Permalink
Replace Model._get_meta_data with hash and range key attributes (#639)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpinner-lyft @garrettheel
jpinner-lyft authored and garrettheel committed Jun 10, 2019
1 parent d8af659 commit 1b7eeb0
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 143 deletions.
119 changes: 37 additions & 82 deletions pynamodb/models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,13 @@
import json
import time
import six
import copy
import logging
import warnings
from inspect import getmembers

from six import add_metaclass
from pynamodb.exceptions import DoesNotExist, TableDoesNotExist, TableError
from pynamodb.attributes import Attribute, AttributeContainer, AttributeContainerMeta, MapAttribute
from pynamodb.connection.base import MetaTable
from pynamodb.connection.table import TableConnection
from pynamodb.connection.util import pythonic
from pynamodb.types import HASH, RANGE
Expand Down Expand Up @@ -168,6 +166,11 @@ class MetaModel(AttributeContainerMeta):
"""
def __init__(cls, name, bases, attrs):
super(MetaModel, cls).__init__(name, bases, attrs)
for attr_name, attribute in cls.get_attributes().items():
if attribute.is_hash_key:
cls._hash_keyname = attr_name
if attribute.is_range_key:
cls._range_keyname = attr_name
if isinstance(attrs, dict):
for attr_name, attr_obj in attrs.items():
if attr_name == META_CLASS_NAME:
Expand Down Expand Up @@ -224,7 +227,8 @@ class Model(AttributeContainer):

# These attributes are named to avoid colliding with user defined
# DynamoDB attributes
_meta_table = None
_hash_keyname = None
_range_keyname = None
_indexes = None
_connection = None
_index_classes = None
Expand All @@ -237,14 +241,13 @@ def __init__(self, hash_key=None, range_key=None, **attributes):
:param attrs: A dictionary of attributes to set on this object.
"""
if hash_key is not None:
attributes[self._dynamo_to_python_attr(self._get_meta_data().hash_keyname)] = hash_key
attributes[self._hash_keyname] = hash_key
if range_key is not None:
range_keyname = self._get_meta_data().range_keyname
if range_keyname is None:
if self._range_keyname is None:
raise ValueError(
"This table has no range key, but a range key value was provided: {}".format(range_key)
)
attributes[self._dynamo_to_python_attr(range_keyname)] = range_key
attributes[self._range_keyname] = range_key
super(Model, self).__init__(**attributes)

@classmethod
Expand All @@ -256,8 +259,8 @@ def batch_get(cls, items, consistent_read=None, attributes_to_get=None):
tuples if range keys are used.
"""
items = list(items)
hash_keyname = cls._get_meta_data().hash_keyname
range_keyname = cls._get_meta_data().range_keyname
hash_key_attribute = cls._hash_key_attribute()
range_key_attribute = cls._range_key_attribute()
keys_to_get = []
while items:
if len(keys_to_get) == BATCH_GET_PAGE_LIMIT:
Expand All @@ -274,16 +277,16 @@ def batch_get(cls, items, consistent_read=None, attributes_to_get=None):
else:
keys_to_get = []
item = items.pop()
if range_keyname:
if range_key_attribute:
hash_key, range_key = cls._serialize_keys(item[0], item[1])
keys_to_get.append({
hash_keyname: hash_key,
range_keyname: range_key
hash_key_attribute.attr_name: hash_key,
range_key_attribute.attr_name: range_key
})
else:
hash_key = cls._serialize_keys(item)[0]
keys_to_get.append({
hash_keyname: hash_key
hash_key_attribute.attr_name: hash_key
})

while keys_to_get:
Expand Down Expand Up @@ -315,7 +318,7 @@ def batch_write(cls, auto_commit=True):
def __repr__(self):
if self.Meta.table_name:
serialized = self._serialize(null_check=False)
if self._get_meta_data().range_keyname:
if self._range_keyname:
msg = "{}<{}, {}>".format(self.Meta.table_name, serialized.get(HASH), serialized.get(RANGE))
else:
msg = "{}<{}>".format(self.Meta.table_name, serialized.get(HASH))
Expand Down Expand Up @@ -412,30 +415,16 @@ def from_raw_data(cls, data):
:param data: A serialized DynamoDB object
"""
mutable_data = copy.copy(data)
if mutable_data is None:
raise ValueError("Received no mutable_data to construct object")
hash_keyname = cls._get_meta_data().hash_keyname
range_keyname = cls._get_meta_data().range_keyname
hash_key_type = cls._get_meta_data().get_attribute_type(hash_keyname)
hash_key = mutable_data.pop(hash_keyname).get(hash_key_type)

hash_key_attr = cls.get_attributes().get(cls._dynamo_to_python_attr(hash_keyname))
if data is None:
raise ValueError("Received no data to construct object")

hash_key = hash_key_attr.deserialize(hash_key)
args = (hash_key,)
kwargs = {}
if range_keyname:
range_key_attr = cls.get_attributes().get(cls._dynamo_to_python_attr(range_keyname))
range_key_type = cls._get_meta_data().get_attribute_type(range_keyname)
range_key = mutable_data.pop(range_keyname).get(range_key_type)
kwargs['range_key'] = range_key_attr.deserialize(range_key)
for name, value in mutable_data.items():
attributes = {}
for name, value in data.items():
attr_name = cls._dynamo_to_python_attr(name)
attr = cls.get_attributes().get(attr_name, None)
if attr:
kwargs[attr_name] = attr.deserialize(attr.get_value(value))
return cls(*args, **kwargs)
attributes[attr_name] = attr.deserialize(attr.get_value(value))
return cls(**attributes)

@classmethod
def count(cls,
Expand Down Expand Up @@ -464,17 +453,8 @@ def count(cls,
cls._get_indexes()
if index_name:
hash_key = cls._index_classes[index_name]._hash_key_attribute().serialize(hash_key)
key_attribute_classes = cls._index_classes[index_name]._get_attributes()
non_key_attribute_classes = cls.get_attributes()
else:
hash_key = cls._serialize_keys(hash_key)[0]
non_key_attribute_classes = dict(cls.get_attributes())
key_attribute_classes = dict(cls.get_attributes())
for name, attr in cls.get_attributes().items():
if attr.is_range_key or attr.is_hash_key:
key_attribute_classes[name] = attr
else:
non_key_attribute_classes[name] = attr

query_args = (hash_key,)
query_kwargs = dict(
Expand Down Expand Up @@ -531,17 +511,8 @@ def query(cls,
cls._get_indexes()
if index_name:
hash_key = cls._index_classes[index_name]._hash_key_attribute().serialize(hash_key)
key_attribute_classes = cls._index_classes[index_name]._get_attributes()
non_key_attribute_classes = cls.get_attributes()
else:
hash_key = cls._serialize_keys(hash_key)[0]
non_key_attribute_classes = {}
key_attribute_classes = {}
for name, attr in cls.get_attributes().items():
if attr.is_range_key or attr.is_hash_key:
key_attribute_classes[name] = attr
else:
non_key_attribute_classes[name] = attr

if page_size is None:
page_size = limit
Expand Down Expand Up @@ -730,14 +701,16 @@ def _from_data(cls, data):
hash_key, attrs = data
range_key = attrs.pop('range_key', None)
attributes = attrs.pop(pythonic(ATTRIBUTES))
hash_keyname = cls._get_meta_data().hash_keyname
hash_keytype = cls._get_meta_data().get_attribute_type(hash_keyname)
hash_key_attribute = cls._hash_key_attribute()
hash_keyname = hash_key_attribute.attr_name
hash_keytype = ATTR_TYPE_MAP[hash_key_attribute.attr_type]
attributes[hash_keyname] = {
hash_keytype: hash_key
}
if range_key is not None:
range_keyname = cls._get_meta_data().range_keyname
range_keytype = cls._get_meta_data().get_attribute_type(range_keyname)
range_key_attribute = cls._range_key_attribute()
range_keyname = range_key_attribute.attr_name
range_keytype = ATTR_TYPE_MAP[range_key_attribute.attr_type]
attributes[range_keyname] = {
range_keytype: range_key
}
Expand Down Expand Up @@ -843,26 +816,18 @@ def _get_save_args(self, attributes=True, null_check=True):
return args, kwargs

@classmethod
def _range_key_attribute(cls):
def _hash_key_attribute(cls):
"""
Returns the attribute class for the hash key
"""
attributes = cls.get_attributes()
range_keyname = cls._get_meta_data().range_keyname
if range_keyname:
attr = attributes[cls._dynamo_to_python_attr(range_keyname)]
else:
attr = None
return attr
return cls.get_attributes()[cls._hash_keyname]

@classmethod
def _hash_key_attribute(cls):
def _range_key_attribute(cls):
"""
Returns the attribute class for the hash key
Returns the attribute class for the range key
"""
attributes = cls.get_attributes()
hash_keyname = cls._get_meta_data().hash_keyname
return attributes[cls._dynamo_to_python_attr(hash_keyname)]
return cls.get_attributes()[cls._range_keyname] if cls._range_keyname else None

def _get_keys(self):
"""
Expand All @@ -871,12 +836,11 @@ def _get_keys(self):
serialized = self._serialize(null_check=False)
hash_key = serialized.get(HASH)
range_key = serialized.get(RANGE, None)
hash_keyname = self._get_meta_data().hash_keyname
range_keyname = self._get_meta_data().range_keyname
attrs = {
hash_keyname: hash_key,
self._hash_key_attribute().attr_name: hash_key,
}
if range_keyname is not None:
if self._range_keyname is not None:
range_keyname = self._range_key_attribute().attr_name
attrs[range_keyname] = range_key
return attrs

Expand All @@ -898,15 +862,6 @@ def _batch_get_page(cls, keys_to_get, consistent_read, attributes_to_get):
unprocessed_items = data.get(UNPROCESSED_KEYS).get(cls.Meta.table_name, {}).get(KEYS, None)
return item_data, unprocessed_items

@classmethod
def _get_meta_data(cls):
"""
A helper object that contains meta data about this table
"""
if cls._meta_table is None:
cls._meta_table = MetaTable(cls._get_connection().describe_table())
return cls._meta_table

@classmethod
def _get_connection(cls):
"""
Expand Down
2 changes: 2 additions & 0 deletions pynamodb/models.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ KeyType = Union[Text, bytes, float, int, Tuple]
class Model(metaclass=MetaModel):
DoesNotExist = DoesNotExist
attribute_values: Dict[Text, Any]
_hash_keyname: Optional[str]
_range_keyname: Optional[str]
_connection: Optional[TableConnection]
def __init__(self, hash_key: Optional[KeyType] = ..., range_key: Optional[Any] = ..., **attrs) -> None: ...
@classmethod
Expand Down

0 comments on commit 1b7eeb0

Please sign in to comment.