Coverage for dynamodx / transact_get.py: 98%

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(self, kset: PrimaryKeySet, *, flatten_top: bool = True): 

29 table_name = self._table_name 

30 transact_items: list[TransactGetItemTypeDef] = [ 

31 _build_get(pk, table_name) for pk in kset.pairs 

32 ] 

33 

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

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

36 

37 if flatten_top and items: 

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

39 else: 

40 head, tail = {}, items 

41 

42 pairs = kset.pairs[1:] if flatten_top else kset.pairs 

43 nested = { 

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

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

46 if item 

47 } 

48 return {**head, **nested} 

49 

50 

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

52 sk = pk[pk.name_sk] 

53 attrs: GetTypeDef = { 

54 'TableName': pk.table_name or table_name, 

55 'Key': serialize(pk), 

56 } 

57 

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

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

60 

61 if projection_expr is not None: 

62 attrs['ProjectionExpression'] = projection_expr 

63 

64 if expr_attr_names is not None: 

65 attrs['ExpressionAttributeNames'] = expr_attr_names 

66 

67 return {'Get': attrs} 

68 

69 

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

71 sk = pk.sk 

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

73 

74 

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

76 sk = pk.sk 

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

78 

79 if path_spec is None: 

80 return item 

81 

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