DSS Wrappers: DynamoDB update

dss/collections/owner_lookup.py

@@ -20,10 +20,10 @@ def put_collection(owner: str, collection_fqid: str, permission_level: str = 'ow
 
 
 def get_collection(owner: str, collection_fqid: str):
-    return dynamodb.get_item(table=collection_db_table,
-                             hash_key=owner,
-                             sort_key=collection_fqid,
-                             return_key='sort_key')
+    collection = dynamodb.get_item(table=collection_db_table,
+                                   hash_key=owner,
+                                   sort_key=collection_fqid)
+    return collection.get('sort_key')
 
 
 def get_collection_fqids_for_owner(owner: str):

dss/dynamodb/__init__.py

@@ -1,3 +1,4 @@
+import typing
 from typing import Generator, Optional
 
 from dss.util.aws.clients import dynamodb as db  # type: ignore
@@ -6,6 +7,9 @@
 class DynamoDBItemNotFound(Exception):
     pass
 
+class SchemaMissMatch(Exception):
+    pass
+
 
 def _format_item(hash_key: str, sort_key: Optional[str], value: Optional[str], ttl: Optional[int] = None) -> dict:
     item = {'hash_key': {'S': hash_key}}
@@ -43,26 +47,26 @@ def put_item(*, table: str, hash_key: str, sort_key: Optional[str] = None, value
     db.put_item(**query)
 
 
-def get_item(*, table: str, hash_key: str, sort_key: Optional[str] = None, return_key: str = 'body') -> str:
+def get_item(*, table: str, hash_key: str, sort_key: str = None) -> typing.Dict:
     """
     Get associated value for a given set of keys from a dynamoDB table.
-
-    Will determine the type of db this is being called on by the number of keys provided (omit
-    sort_key to GET a value from a db with only 1 primary key).
-
     :param table: Name of the table in AWS.
     :param str hash_key: 1st primary key that can be used to fetch associated sort_keys and values.
     :param str sort_key: 2nd primary key, used with hash_key to fetch a specific value.
-                         Note: If not specified, this will GET only 1 key (hash_key) and 1 value.
-    :param str return_key: Either "body" (to return all values) or "sort_key" (to return all 2nd primary keys).
-    :return: None or str
+    :return: item object from ddb (dic)
     """
+    return_value = {}
     query = {'TableName': table,
              'Key': _format_item(hash_key=hash_key, sort_key=sort_key, value=None)}
-    item = db.get_item(**query).get('Item')
+    try:
+        item = db.get_item(**query).get('Item')
+    except db.exceptions.ValidationException:
+        raise SchemaMissMatch(f'Query schema {query} does not match table {table}')
     if item is None:
         raise DynamoDBItemNotFound(f'Query failed to fetch item from database: {query}')
-    return item[return_key]['S']
+    for k, v in item.items():  # strips out ddb typing info
+        return_value[k] = [*v.values()][0]
+    return return_value
 
 
 def get_primary_key_items(*, table: str, key: str, return_key: str = 'body') -> Generator[str, None, None]:
@@ -138,3 +142,27 @@ def delete_item(*, table: str, hash_key: str, sort_key: Optional[str] = None):
     query = {'TableName': table,
              'Key': _format_item(hash_key=hash_key, sort_key=sort_key, value=None)}
     db.delete_item(**query)
+
+
+def update_item(*, table: str, hash_key: str, sort_key: Optional[str] = None, update_expression: Optional[str],
+                expression_attribute_values: typing.Dict):
+    """
+    Update an item from a dynamoDB table.
+    Will determine the type of db this is being called on by the number of keys provided (omit
+    sort_key to UPDATE from a db with only 1 primary key).
+    NOTE:
+    https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_UpdateItem.html
+    :param table: Name of the table in AWS.
+    :param str hash_key: 1st primary key that can be used to fetch associated sort_keys and values.
+    :param str sort_key: 2nd primary key, used with hash_key to fetch a specific value.
+                         Note: If not specified, this will DELETE only 1 key (hash_key) and 1 value.
+    :param str update_expression: Expression used to update value, needs action to be performed and new value
+    :param str expression_attribute_values: attribute values to use from the expression
+    :return: None
+    """
+    query = {'TableName': table,
+             'Key': _format_item(hash_key=hash_key, sort_key=sort_key, value=None)}
+    if update_expression:
+        query['UpdateExpression'] = update_expression
+        query['ExpressionAttributeValues'] = expression_attribute_values
+    db.update_item(**query)

dss/subscriptions_v2/__init__.py

@@ -33,7 +33,8 @@ def get_subscription(replica: Replica, owner: str, uuid: str):
         item = dynamodb.get_item(table=subscription_db_table.format(replica.name),
                                  hash_key=owner,
                                  sort_key=uuid)
-        return json.loads(item)
+
+        return json.loads(item.get('body'))
     except dynamodb.DynamoDBItemNotFound:
         return None
 

dss/util/async_state.py

@@ -47,7 +47,7 @@ def put(cls, key: str, body: dict = None) -> typing.Any:
     def get(cls, key: str) -> typing.Any:
         try:
             item = get_item(table=cls.table, hash_key=key)
-            body = json.loads(item)
+            body = json.loads(item.get('body'))
             item_class = _all_subclasses(cls)[body['_type']]
             return item_class(key, body)
         except DynamoDBItemNotFound: