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

144 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-03-31 13:12 -0600

1"""Utility helper functions""" 

2 

3import sys 

4import typing as t 

5import random 

6import string 

7import logging 

8import datetime 

9from pathlib import Path 

10from pprint import pformat 

11from shutil import rmtree 

12from tempfile import mkdtemp 

13from git import Repo 

14import tiered_debug as debug 

15from ..defaults import ilm_force_merge, ilm_phase, TIER 

16from ..exceptions import TestbedMisconfig 

17 

18logger = logging.getLogger(__name__) 

19 

20 

21def build_ilm_phase( 

22 phase: str, 

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

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

25 fm: bool = False, 

26) -> t.Dict: 

27 """Build a single ILM policy step based on phase""" 

28 debug.lv2('Starting function...') 

29 retval = ilm_phase(phase) 

30 if phase in ['cold', 'frozen']: 

31 if repo: 

32 retval[phase]['actions']['searchable_snapshot'] = { 

33 'snapshot_repository': repo, 

34 'force_merge_index': fm, 

35 } 

36 else: 

37 msg = ( 

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

39 f'provided' 

40 ) 

41 raise TestbedMisconfig(msg) 

42 if actions: 

43 retval[phase]['actions'].update(actions) 

44 debug.lv3('Exiting function, returing value') 

45 debug.lv5(f'Value = {retval}') 

46 return retval 

47 

48 

49def build_ilm_policy( 

50 phases: list = None, 

51 forcemerge: bool = False, 

52 max_num_segments: int = 1, 

53 readonly: t.Union[str, None] = None, 

54 repository: t.Union[str, None] = None, 

55) -> t.Dict: 

56 """ 

57 Build a full ILM policy based on the provided phases. 

58 Put forcemerge in the last phase before cold or frozen (whichever comes first) 

59 """ 

60 debug.lv2('Starting function...') 

61 if not phases: 

62 phases = ['hot', 'delete'] 

63 retval = {} 

64 if ('cold' in phases or 'frozen' in phases) and not repository: 

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

66 for phase in phases: 

67 actions = None 

68 if readonly == phase: 

69 actions = {"readonly": {}} 

70 phase = build_ilm_phase(phase, repo=repository, fm=forcemerge, actions=actions) 

71 retval.update(phase) 

72 if forcemerge: 

73 retval['hot']['actions'].update( 

74 ilm_force_merge(max_num_segments=max_num_segments) 

75 ) 

76 debug.lv3('Exiting function, returing value') 

77 debug.lv5("Value = {'phases': " + f'{retval}' + "}") 

78 return {'phases': retval} 

79 

80 

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

82 """Return the routing allocation tier preference""" 

83 debug.lv2('Starting function...') 

84 try: 

85 debug.lv4(f'Checking for tier: {tier}') 

86 pref = TIER[tier]['pref'] 

87 except KeyError: 

88 # Fallback value 

89 pref = 'data_content' 

90 retval = {'index.routing.allocation.include._tier_preference': pref} 

91 debug.lv3('Exiting function, returing value') 

92 debug.lv5(f'Value = {retval}') 

93 return retval 

94 

95 

96def iso8601_now() -> str: 

97 """ 

98 :returns: An ISO8601 timestamp based on now 

99 :rtype: str 

100 """ 

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

102 # 

103 # ## Example: 

104 # ## The new way: 

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

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

107 # ## End Example 

108 # 

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

110 # UTC timezone 

111 # 

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

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

114 # Elasticsearch uses) 

115 # 

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

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

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

119 if len(parts) == 1: 

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

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

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

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

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

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

126 

127 

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

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

130 debug.lv2('Starting function...') 

131 retval = f'{TIER[tier]["prefix"]}-{index}' 

132 debug.lv3('Exiting function, returing value') 

133 debug.lv5(f'Value = {retval}') 

134 return retval 

135 

136 

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

138 """ 

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

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

141 logging and showing more readable dictionaries. 

142 

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

144 compact, sort_dicts and underscore_numbers are passed to the PrettyPrinter 

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

146 

147 The keyword arg, ``underscore_numbers`` is only available in Python versions 

148 3.10 and up, so there is a test here to add it when that is the case. 

149 """ 

150 defaults = [ 

151 ('indent', 2), 

152 ('width', 80), 

153 ('depth', None), 

154 ('compact', False), 

155 ('sort_dicts', False), 

156 ] 

157 vinfo = python_version() 

158 if vinfo[0] == 3 and vinfo[1] >= 10: 

159 # underscore_numbers only works in 3.10 and up 

160 defaults.append(('underscore_numbers', False)) 

161 kw = {} 

162 for tup in defaults: 

163 key, default = tup 

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

165 

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

167 

168 

169def process_preset( 

170 builtin: t.Union[str, None], 

171 path: t.Union[str, None], 

172 ref: t.Union[str, None], 

173 url: t.Union[str, None], 

174) -> t.Tuple: 

175 """Process the preset settings 

176 :param preset: One of `builtin`, `git`, or `path` 

177 :param builtin: The name of a builtin preset 

178 :param path: A relative or absolute file path. Used by presets `git` and `path` 

179 :param ref: A Git ref (e.g. 'main'). Only used by preset `git` 

180 :param url: A Git repository URL. Only used by preset `git` 

181 """ 

182 debug.lv2('Starting function...') 

183 modpath = None 

184 tmpdir = None 

185 if builtin: # Overrides any other options 

186 modpath = f'es_testbed.presets.{builtin}' 

187 else: 

188 trygit = False 

189 try: 

190 debug.lv4('TRY: Checking for git preset') 

191 kw = {'path': path, 'ref': ref, 'url': url} 

192 raise_on_none(**kw) 

193 trygit = True # We have all 3 kwargs necessary for git 

194 except ValueError as resp: # Not able to do a git preset 

195 debug.lv1(f'Unable to import a git-based preset: {resp}') 

196 if trygit: # Trying a git import 

197 tmpdir = mkdtemp() 

198 try: 

199 debug.lv4('TRY: Attempting to clone git repository') 

200 _ = Repo.clone_from(url, tmpdir, branch=ref, depth=1) 

201 filepath = Path(tmpdir) / path 

202 except Exception as err: 

203 logger.error(f'Git clone failed: {err}') 

204 rmtree(tmpdir) # Clean up after failed attempt 

205 raise err 

206 filepath = Path(path) # It should work even if path is None 

207 if not filepath.resolve().is_dir(): 

208 raise ValueError(f'The provided path "{path}" is not a directory') 

209 modpath = filepath.resolve().name # The final dirname 

210 parent = filepath.parent.resolve() # Up one level 

211 # We now make the parent path part of the sys.path. 

212 sys.path.insert(0, parent) # This should persist beyond this module 

213 debug.lv3('Exiting function, returing value') 

214 debug.lv5(f'Value = ({modpath}, {tmpdir})') 

215 return modpath, tmpdir 

216 

217 

218def python_version() -> t.Tuple: 

219 """ 

220 Return running Python version tuple, e.g. 3.12.2 would be (3, 12, 2) 

221 """ 

222 debug.lv2('Starting function...') 

223 _ = sys.version_info 

224 retval = (_[0], _[1], _[2]) 

225 debug.lv3('Exiting function, returing value') 

226 debug.lv5(f'Value = {retval}') 

227 return retval 

228 

229 

230def raise_on_none(**kwargs): 

231 """Raise if any kwargs have a None value""" 

232 debug.lv2('Starting function...') 

233 for key, value in kwargs.items(): 

234 if value is None: 

235 debug.lv3('Exiting function, raising ValueError') 

236 raise ValueError(f'kwarg "{key}" cannot have a None value') 

237 debug.lv3('Exiting function') 

238 

239 

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

241 """Generate a random string""" 

242 debug.lv2('Starting function...') 

243 letters = string.ascii_uppercase 

244 if lowercase: 

245 letters = string.ascii_lowercase 

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

247 debug.lv3('Exiting function, returing value') 

248 debug.lv5(f'Value = {retval}') 

249 return retval 

250 

251 

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

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

254 debug.lv2('Starting function...') 

255 retval = TIER[tier]["storage"] 

256 debug.lv3('Exiting function, returing value') 

257 debug.lv5(f'Value = {retval}') 

258 return retval