Make sure if a sort_key is provided in the case of dynamodb backend

romaryd · May 30, 2018 · b818387 · b818387
1 parent 0f110b1
commit b818387
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/jsonrepo/__init__.py b/jsonrepo/__init__.py
@@ -5,5 +5,5 @@
 Copyright (C) 2017 Romary Dupuis
 """
 
-__version__ = '0.1.5'
+__version__ = '0.1.6'
 __all__ = ['repository', 'mixin', 'record']
diff --git a/jsonrepo/backends/dynamodb.py b/jsonrepo/backends/dynamodb.py
@@ -25,10 +25,14 @@ def dynamodb_server(self):
 
     def get(self, key, sort_key):
         self.logger.debug('Storage - get {}'.format(self.prefixed(key)))
-        res = self.dynamodb_server.get_item(Key={
+        query = {
             self._key: self.prefixed(key),
-            self._sort_key: sort_key
-        })
+        }
+        if sort_key is not None:
+            query.update({
+                self._sort_key: sort_key
+            })
+        res = self.dynamodb_server.get_item(Key=query)
         if 'Item' in res and 'value' in res['Item']:
             return res['Item']['value']
 
@@ -54,10 +58,14 @@ def set(self, key, sort_key, value):
 
     def delete(self, key, sort_key):
         self.logger.debug('Storage - delete {}'.format(self.prefixed(key)))
-        return self.dynamodb_server.delete_item(Key={
+        query = {
             self._key: self.prefixed(key),
-            self._sort_key: sort_key
-        })
+        }
+        if sort_key is not None:
+            query.update({
+                self._sort_key: sort_key
+            })
+        return self.dynamodb_server.delete_item(Key=query)
 
     def history(self, key, _from='-', _to='+', _desc=True):
         if _from != '-':