Coverage for dynamodx / transact_get.py: 100%

1from typing import TYPE_CHECKING, Any 

2 

3import jmespath 

4 

5from dynamodx.keys import PrimaryKey, PrimaryKeySet 

6from dynamodx.types import deserialize, serialize 

7 

8if TYPE_CHECKING: 

9 from mypy_boto3_dynamodb.client import DynamoDBClient 

10 from mypy_boto3_dynamodb.type_defs import GetTypeDef, TransactGetItemTypeDef 

11 

12else: 

13 DynamoDBClient = Any 

14 GetTypeDef = Any 

15 TransactGetItemTypeDef = Any 

16 

17 

18class TransactGet: 

19 def __init__( 

20 self, 

21 table_name: str, 

22 *, 

23 client: DynamoDBClient, 

24 ) -> None: 

25 self._table_name = table_name 

26 self._client = client 

27 

28 def get_items( 

29 self, 

30 keyset: PrimaryKeySet, 

31 *, 

32 flatten_top: bool = True, 

33 ) -> dict[str, Any]: 

34 """Get multiple items via a transaction based on the provided TransactKey. 

35 

36 Parameters 

37 ---------- 

38 key : PrimaryKeySet 

39 

40 flatten_top : bool, optional 

41 Determines whether the first nested item in the transaction result 

42 should be flattened, 

43 i.e., extracted to serve as the primary item at the top level of 

44 the returned dict. 

45 If True, the nested item is promoted to the top level. 

46 

47 Returns 

48 ------- 

49 dict[str, Any] 

50 A dict of items retrieved from the transaction. 

51 """ 

52 table_name = self._table_name 

53 transact_items: list[TransactGetItemTypeDef] = [ 

54 _build_get(pk, table_name) for pk in keyset.pairs 

55 ] 

56 

57 output = self._client.transact_get_items(TransactItems=transact_items) 

58 items = [deserialize(r.get('Item', {})) for r in output.get('Responses', [])] 

59 

60 if flatten_top and items: 

61 head, tail = items[0], items[1:] 

62 else: 

63 head, tail = {}, items 

64 

65 pairs = keyset.pairs[1:] if flatten_top else keyset.pairs 

66 nested = { 

67 _output_key(pk): _project_item(pk, item) 

68 for pk, item in zip(pairs, tail, strict=True) 

69 if item 

70 } 

71 

72 return {**head, **nested} 

73 

74 

75def _build_get(pk: PrimaryKey, table_name: str) -> TransactGetItemTypeDef: 

76 sk = pk[pk.name_sk] 

77 attrs: GetTypeDef = { 

78 'TableName': pk.table_name or table_name, 

79 'Key': serialize(pk), 

80 } 

81 

82 projection_expr = getattr(sk, 'projection_expr', None) 

83 expr_attr_names = getattr(sk, 'expr_attr_names', None) 

84 

85 if projection_expr is not None: 

86 attrs['ProjectionExpression'] = projection_expr 

87 

88 if expr_attr_names is not None: 

89 attrs['ExpressionAttributeNames'] = expr_attr_names 

90 

91 return {'Get': attrs} 

92 

93 

94def _output_key(pk: PrimaryKey) -> str: 

95 sk = pk.sk 

96 return str(getattr(sk, 'rename_key', sk)) 

97 

98 

99def _project_item(pk: PrimaryKey, item: dict) -> dict: 

100 sk = pk.sk 

101 path_spec = getattr(sk, 'path_spec', None) 

102 

103 if path_spec is None: 

104 return item 

105 

106 return jmespath.compile(path_spec).search(item)