1from boto.dynamodb2.types import STRING
2
3
4class BaseSchemaField(object):
5    """
6    An abstract class for defining schema fields.
7
8    Contains most of the core functionality for the field. Subclasses must
9    define an ``attr_type`` to pass to DynamoDB.
10    """
11    attr_type = None
12
13    def __init__(self, name, data_type=STRING):
14        """
15        Creates a Python schema field, to represent the data to pass to
16        DynamoDB.
17
18        Requires a ``name`` parameter, which should be a string name of the
19        field.
20
21        Optionally accepts a ``data_type`` parameter, which should be a
22        constant from ``boto.dynamodb2.types``. (Default: ``STRING``)
23        """
24        self.name = name
25        self.data_type = data_type
26
27    def definition(self):
28        """
29        Returns the attribute definition structure DynamoDB expects.
30
31        Example::
32
33            >>> field.definition()
34            {
35                'AttributeName': 'username',
36                'AttributeType': 'S',
37            }
38
39        """
40        return {
41            'AttributeName': self.name,
42            'AttributeType': self.data_type,
43        }
44
45    def schema(self):
46        """
47        Returns the schema structure DynamoDB expects.
48
49        Example::
50
51            >>> field.schema()
52            {
53                'AttributeName': 'username',
54                'KeyType': 'HASH',
55            }
56
57        """
58        return {
59            'AttributeName': self.name,
60            'KeyType': self.attr_type,
61        }
62
63
64class HashKey(BaseSchemaField):
65    """
66    An field representing a hash key.
67
68    Example::
69
70        >>> from boto.dynamodb2.types import NUMBER
71        >>> HashKey('username')
72        >>> HashKey('date_joined', data_type=NUMBER)
73
74    """
75    attr_type = 'HASH'
76
77
78class RangeKey(BaseSchemaField):
79    """
80    An field representing a range key.
81
82    Example::
83
84        >>> from boto.dynamodb2.types import NUMBER
85        >>> HashKey('username')
86        >>> HashKey('date_joined', data_type=NUMBER)
87
88    """
89    attr_type = 'RANGE'
90
91
92class BaseIndexField(object):
93    """
94    An abstract class for defining schema indexes.
95
96    Contains most of the core functionality for the index. Subclasses must
97    define a ``projection_type`` to pass to DynamoDB.
98    """
99    def __init__(self, name, parts):
100        self.name = name
101        self.parts = parts
102
103    def definition(self):
104        """
105        Returns the attribute definition structure DynamoDB expects.
106
107        Example::
108
109            >>> index.definition()
110            {
111                'AttributeName': 'username',
112                'AttributeType': 'S',
113            }
114
115        """
116        definition = []
117
118        for part in self.parts:
119            definition.append({
120                'AttributeName': part.name,
121                'AttributeType': part.data_type,
122            })
123
124        return definition
125
126    def schema(self):
127        """
128        Returns the schema structure DynamoDB expects.
129
130        Example::
131
132            >>> index.schema()
133            {
134                'IndexName': 'LastNameIndex',
135                'KeySchema': [
136                    {
137                        'AttributeName': 'username',
138                        'KeyType': 'HASH',
139                    },
140                ],
141                'Projection': {
142                    'ProjectionType': 'KEYS_ONLY',
143                }
144            }
145
146        """
147        key_schema = []
148
149        for part in self.parts:
150            key_schema.append(part.schema())
151
152        return {
153            'IndexName': self.name,
154            'KeySchema': key_schema,
155            'Projection': {
156                'ProjectionType': self.projection_type,
157            }
158        }
159
160
161class AllIndex(BaseIndexField):
162    """
163    An index signifying all fields should be in the index.
164
165    Example::
166
167        >>> AllIndex('MostRecentlyJoined', parts=[
168        ...     HashKey('username'),
169        ...     RangeKey('date_joined')
170        ... ])
171
172    """
173    projection_type = 'ALL'
174
175
176class KeysOnlyIndex(BaseIndexField):
177    """
178    An index signifying only key fields should be in the index.
179
180    Example::
181
182        >>> KeysOnlyIndex('MostRecentlyJoined', parts=[
183        ...     HashKey('username'),
184        ...     RangeKey('date_joined')
185        ... ])
186
187    """
188    projection_type = 'KEYS_ONLY'
189
190
191class IncludeIndex(BaseIndexField):
192    """
193    An index signifying only certain fields should be in the index.
194
195    Example::
196
197        >>> IncludeIndex('GenderIndex', parts=[
198        ...     HashKey('username'),
199        ...     RangeKey('date_joined')
200        ... ], includes=['gender'])
201
202    """
203    projection_type = 'INCLUDE'
204
205    def __init__(self, *args, **kwargs):
206        self.includes_fields = kwargs.pop('includes', [])
207        super(IncludeIndex, self).__init__(*args, **kwargs)
208
209    def schema(self):
210        schema_data = super(IncludeIndex, self).schema()
211        schema_data['Projection']['NonKeyAttributes'] = self.includes_fields
212        return schema_data
213
214
215class GlobalBaseIndexField(BaseIndexField):
216    """
217    An abstract class for defining global indexes.
218
219    Contains most of the core functionality for the index. Subclasses must
220    define a ``projection_type`` to pass to DynamoDB.
221    """
222    throughput = {
223        'read': 5,
224        'write': 5,
225    }
226
227    def __init__(self, *args, **kwargs):
228        throughput = kwargs.pop('throughput', None)
229
230        if throughput is not None:
231            self.throughput = throughput
232
233        super(GlobalBaseIndexField, self).__init__(*args, **kwargs)
234
235    def schema(self):
236        """
237        Returns the schema structure DynamoDB expects.
238
239        Example::
240
241            >>> index.schema()
242            {
243                'IndexName': 'LastNameIndex',
244                'KeySchema': [
245                    {
246                        'AttributeName': 'username',
247                        'KeyType': 'HASH',
248                    },
249                ],
250                'Projection': {
251                    'ProjectionType': 'KEYS_ONLY',
252                },
253                'ProvisionedThroughput': {
254                    'ReadCapacityUnits': 5,
255                    'WriteCapacityUnits': 5
256                }
257            }
258
259        """
260        schema_data = super(GlobalBaseIndexField, self).schema()
261        schema_data['ProvisionedThroughput'] = {
262            'ReadCapacityUnits': int(self.throughput['read']),
263            'WriteCapacityUnits': int(self.throughput['write']),
264        }
265        return schema_data
266
267
268class GlobalAllIndex(GlobalBaseIndexField):
269    """
270    An index signifying all fields should be in the index.
271
272    Example::
273
274        >>> GlobalAllIndex('MostRecentlyJoined', parts=[
275        ...     HashKey('username'),
276        ...     RangeKey('date_joined')
277        ... ],
278        ... throughput={
279        ...     'read': 2,
280        ...     'write': 1,
281        ... })
282
283    """
284    projection_type = 'ALL'
285
286
287class GlobalKeysOnlyIndex(GlobalBaseIndexField):
288    """
289    An index signifying only key fields should be in the index.
290
291    Example::
292
293        >>> GlobalKeysOnlyIndex('MostRecentlyJoined', parts=[
294        ...     HashKey('username'),
295        ...     RangeKey('date_joined')
296        ... ],
297        ... throughput={
298        ...     'read': 2,
299        ...     'write': 1,
300        ... })
301
302    """
303    projection_type = 'KEYS_ONLY'
304
305
306class GlobalIncludeIndex(GlobalBaseIndexField, IncludeIndex):
307    """
308    An index signifying only certain fields should be in the index.
309
310    Example::
311
312        >>> GlobalIncludeIndex('GenderIndex', parts=[
313        ...     HashKey('username'),
314        ...     RangeKey('date_joined')
315        ... ],
316        ... includes=['gender'],
317        ... throughput={
318        ...     'read': 2,
319        ...     'write': 1,
320        ... })
321
322    """
323    projection_type = 'INCLUDE'
324
325    def __init__(self, *args, **kwargs):
326        throughput = kwargs.pop('throughput', None)
327        IncludeIndex.__init__(self, *args, **kwargs)
328        if throughput:
329            kwargs['throughput'] = throughput
330        GlobalBaseIndexField.__init__(self, *args, **kwargs)
331
332    def schema(self):
333        # Pick up the includes.
334        schema_data = IncludeIndex.schema(self)
335        # Also the throughput.
336        schema_data.update(GlobalBaseIndexField.schema(self))
337        return schema_data
338