Coverage for /Users/buh/.pyenv/versions/3.12.2/envs/es-testbed/lib/python3.12/site-packages/es_testbed/helpers/utils.py: 89%

82 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-05-02 23:00 -0600

1"""Utility helper functions""" 

2 

3import typing as t 

4from pprint import pformat 

5import random 

6import string 

7import logging 

8from datetime import datetime, timezone 

9from ..defaults import ilm_force_merge, ilm_phase, MAPPING, TIER 

10from ..exceptions import TestbedMisconfig 

11 

12logger = logging.getLogger(__name__) 

13 

14 

15def build_ilm_phase( 

16 tier: str, 

17 actions: t.Union[t.Dict, None] = None, 

18 repo: t.Union[str, None] = None, 

19 fm: bool = False, 

20) -> t.Dict: 

21 """Build a single ILM policy step based on tier""" 

22 phase = ilm_phase(tier) 

23 if tier in ['cold', 'frozen']: 

24 if repo: 

25 phase[tier]['actions']['searchable_snapshot'] = { 

26 'snapshot_repository': repo, 

27 'force_merge_index': fm, 

28 } 

29 else: 

30 msg = ( 

31 f'Unable to build ILM phase for {tier} tier. Value for repository not ' 

32 f'provided' 

33 ) 

34 raise TestbedMisconfig(msg) 

35 if actions: 

36 phase[tier]['actions'].update(actions) 

37 return phase 

38 

39 

40def build_ilm_policy( 

41 tiers: list = None, 

42 forcemerge: bool = False, 

43 max_num_segments: int = 1, 

44 repository: str = None, 

45) -> t.Dict: 

46 """ 

47 Build a full ILM policy based on the provided tiers. 

48 Put forcemerge in the last tier before cold or frozen (whichever comes first) 

49 """ 

50 if not tiers: 

51 tiers = ['hot', 'delete'] 

52 phases = {} 

53 if ('cold' in tiers or 'frozen' in tiers) and not repository: 

54 raise TestbedMisconfig('Cannot build cold or frozen phase without repository') 

55 for tier in tiers: 

56 phase = build_ilm_phase(tier, repo=repository, fm=forcemerge) 

57 phases.update(phase) 

58 if forcemerge: 

59 phases['hot']['actions'].update( 

60 ilm_force_merge(max_num_segments=max_num_segments) 

61 ) 

62 return {'phases': phases} 

63 

64 

65def doc_gen( 

66 count: int = 10, start_at: int = 0, match: bool = True 

67) -> t.Generator[t.Dict, None, None]: 

68 """Create this doc for each count""" 

69 keys = ['message', 'nested', 'deep'] 

70 # Start with an empty map 

71 matchmap = {} 

72 # Iterate over each key 

73 for key in keys: 

74 # If match is True 

75 if match: 

76 # Set matchmap[key] to key 

77 matchmap[key] = key 

78 else: 

79 # Otherwise matchmap[key] will have a random string value 

80 matchmap[key] = randomstr() 

81 

82 # This is where count and start_at matter 

83 for num in range(start_at, start_at + count): 

84 yield { 

85 '@timestamp': iso8601_now(), 

86 'message': f'{matchmap["message"]}{num}', # message# or randomstr# 

87 'number': ( 

88 num if match else random.randint(1001, 32767) 

89 ), # value of num or random int 

90 'nested': {'key': f'{matchmap["nested"]}{num}'}, # nested# 

91 'deep': {'l1': {'l2': {'l3': f'{matchmap["deep"]}{num}'}}}, # deep# 

92 } 

93 

94 

95def getlogger(name: str) -> logging.getLogger: 

96 """Return a named logger""" 

97 return logging.getLogger(name) 

98 

99 

100def get_routing(tier='hot') -> t.Dict: 

101 """Return the routing allocation tier preference""" 

102 try: 

103 pref = TIER[tier]['pref'] 

104 except KeyError: 

105 # Fallback value 

106 pref = 'data_content' 

107 return {'index.routing.allocation.include._tier_preference': pref} 

108 

109 

110def iso8601_now() -> str: 

111 """ 

112 :returns: An ISO8601 timestamp based on now 

113 :rtype: str 

114 """ 

115 # Because Python 3.12 now requires non-naive timezone declarations, we must change. 

116 # 

117 # ## Example: 

118 # ## The new way: 

119 # ## datetime.now(timezone.utc).isoformat() 

120 # ## Result: 2024-04-16T16:00:00+00:00 

121 # ## End Example 

122 # 

123 # Note that the +00:00 is appended now where we affirmatively declare the 

124 # UTC timezone 

125 # 

126 # As a result, we will use this function to prune away the timezone if it is 

127 # +00:00 and replace it with Z, which is shorter Zulu notation for UTC (which 

128 # Elasticsearch uses) 

129 # 

130 # We are MANUALLY, FORCEFULLY declaring timezone.utc, so it should ALWAYS be 

131 # +00:00, but could in theory sometime show up as a Z, so we test for that. 

132 

133 parts = datetime.now(timezone.utc).isoformat().split('+') 

134 if len(parts) == 1: 

135 if parts[0][-1] == 'Z': 

136 return parts[0] # Our ISO8601 already ends with a Z for Zulu/UTC time 

137 return f'{parts[0]}Z' # It doesn't end with a Z so we put one there 

138 if parts[1] == '00:00': 

139 return f'{parts[0]}Z' # It doesn't end with a Z so we put one there 

140 return f'{parts[0]}+{parts[1]}' # Fallback publishes the +TZ, whatever that was 

141 

142 

143def mapping_component() -> t.Dict: 

144 """Return a mappings component template""" 

145 return {'mappings': MAPPING} 

146 

147 

148def mounted_name(index: str, tier: str): 

149 """Return a value for renamed_index for mounting a searchable snapshot index""" 

150 return f'{TIER[tier]["prefix"]}-{index}' 

151 

152 

153def prettystr(*args, **kwargs) -> str: 

154 """ 

155 A (nearly) straight up wrapper for pprint.pformat, except that we provide our own 

156 default values for 'indent' (2) and 'sort_dicts' (False). Primarily for debug 

157 logging and showing more readable dictionaries. 

158 

159 'Return the formatted representation of object as a string. indent, width, depth, 

160 compact, sort_dicts and underscore_numbers are passed to the PrettyPrinter 

161 constructor as formatting parameters' (from pprint documentation). 

162 """ 

163 defaults = [ 

164 ('indent', 2), 

165 ('width', 80), 

166 ('depth', None), 

167 ('compact', False), 

168 ('sort_dicts', False), 

169 ('underscore_numbers', False), 

170 ] 

171 kw = {} 

172 for tup in defaults: 

173 key, default = tup 

174 kw[key] = kwargs[key] if key in kwargs else default 

175 

176 return f"\n{pformat(*args, **kw)}" # newline in front so it always starts clean 

177 

178 

179def randomstr(length: int = 16, lowercase: bool = False) -> str: 

180 """Generate a random string""" 

181 letters = string.ascii_uppercase 

182 if lowercase: 

183 letters = string.ascii_lowercase 

184 return str(''.join(random.choices(letters + string.digits, k=length))) 

185 

186 

187def setting_component( 

188 ilm_policy: t.Union[str, None] = None, rollover_alias: t.Union[str, None] = None 

189) -> t.Dict[str, t.Any]: 

190 """Return a settings component template""" 

191 val = {'settings': {'index.number_of_replicas': 0}} 

192 if ilm_policy: 

193 val['settings']['index.lifecycle.name'] = ilm_policy 

194 if rollover_alias: 

195 val['settings']['index.lifecycle.rollover_alias'] = rollover_alias 

196 return val 

197 

198 

199def storage_type(tier: str) -> t.Dict: 

200 """Return the storage type of a searchable snapshot by tier""" 

201 return TIER[tier]["storage"]