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 PrimaryKeySet. 

35 

36 Parameters 

37 ---------- 

38 keyset : PrimaryKeySet 

39 Primary keys of the items to fetch in the transaction. 

40 

41 flatten_top : bool, optional 

42 Determines whether the first nested item in the transaction result 

43 should be flattened, 

44 

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

46 the returned dict. 

47 

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

49 

50 Returns 

51 ------- 

52 dict[str, Any] 

53 A dict of items retrieved from the transaction. 

54 

55 Notes 

56 ----- 

57 Missing items are ignored and are not included in the returned dict. 

58 The order of items follows the order defined in ``keyset.pairs``. 

59 """ 

60 table_name = self._table_name 

61 transact_items: list[TransactGetItemTypeDef] = [ 

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

63 ] 

64 

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

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

67 

68 if flatten_top and items: 

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

70 else: 

71 head, tail = {}, items 

72 

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

74 nested = { 

75 _output_key(pk): project_item(pk, item) 

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

77 if item 

78 } 

79 

80 return {**head, **nested} 

81 

82 

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

84 sk = pk[pk.name_sk] 

85 attrs: GetTypeDef = { 

86 'TableName': pk.table_name or table_name, 

87 'Key': serialize(pk), 

88 } 

89 

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

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

92 

93 if projection_expr is not None: 

94 attrs['ProjectionExpression'] = projection_expr 

95 

96 if expr_attr_names is not None: 

97 attrs['ExpressionAttributeNames'] = expr_attr_names 

98 

99 return {'Get': attrs} 

100 

101 

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

103 sk = pk.sk 

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

105 

106 

107def project_item(pk: PrimaryKey, item: dict) -> dict: 

108 path_spec = getattr(pk.sk, 'path_spec', None) 

109 

110 if path_spec is None: 

111 return item 

112 

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