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

136 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-04-17 22:23 -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 

14from .debug import debug, begin_end 

15from .defaults import ilm_force_merge, ilm_phase, TIER 

16from .exceptions import TestbedMisconfig 

17 

18logger = logging.getLogger(__name__) 

19 

20 

21@begin_end() 

22def build_ilm_phase( 

23 phase: str, 

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

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

26 fm: bool = False, 

27) -> t.Dict: 

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

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.lv5(f'Return value = {retval}') 

45 return retval 

46 

47 

48@begin_end() 

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 if not phases: 

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

62 retval = {} 

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

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

65 for phase in phases: 

66 actions = None 

67 if readonly == phase: 

68 actions = {"readonly": {}} 

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

70 retval.update(phase) 

71 if forcemerge: 

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

73 ilm_force_merge(max_num_segments=max_num_segments) 

74 ) 

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

76 return {'phases': retval} 

77 

78 

79@begin_end() 

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

81 """Return the routing allocation tier preference""" 

82 try: 

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

84 pref = TIER[tier]['pref'] 

85 except KeyError: 

86 # Fallback value 

87 pref = 'data_content' 

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

89 debug.lv5(f'Return value = {retval}') 

90 return retval 

91 

92 

93@begin_end(begin=5, end=5) 

94def iso8601_now() -> str: 

95 """ 

96 :returns: An ISO8601 timestamp based on now 

97 :rtype: str 

98 """ 

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

100 # 

101 # ## Example: 

102 # ## The new way: 

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

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

105 # ## End Example 

106 # 

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

108 # UTC timezone 

109 # 

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

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

112 # Elasticsearch uses) 

113 # 

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

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

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

117 if len(parts) == 1: 

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

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

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

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

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

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

124 

125 

126@begin_end() 

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

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

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

130 debug.lv5(f'Return value = {retval}') 

131 return retval 

132 

133 

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

135 """ 

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

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

138 logging and showing more readable dictionaries. 

139 

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

141 compact, sort_dicts and underscore_numbers are passed to the PrettyPrinter 

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

143 

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

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

146 """ 

147 defaults = [ 

148 ('indent', 2), 

149 ('width', 80), 

150 ('depth', None), 

151 ('compact', False), 

152 ('sort_dicts', False), 

153 ] 

154 vinfo = python_version() 

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

156 # underscore_numbers only works in 3.10 and up 

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

158 kw = {} 

159 for tup in defaults: 

160 key, default = tup 

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

162 

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

164 

165 

166@begin_end() 

167def process_preset( 

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

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

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

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

172) -> t.Tuple: 

173 """Process the preset settings 

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

175 :param builtin: The name of a builtin preset 

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

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

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

179 """ 

180 modpath = None 

181 tmpdir = None 

182 if builtin: # Overrides any other options 

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

184 else: 

185 trygit = False 

186 try: 

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

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

189 raise_on_none(**kw) 

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

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

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

193 if trygit: # Trying a git import 

194 tmpdir = mkdtemp() 

195 try: 

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

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

198 filepath = Path(tmpdir) / path 

199 except Exception as err: 

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

201 rmtree(tmpdir) # Clean up after failed attempt 

202 raise err 

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

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

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

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

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

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

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

210 debug.lv5(f'Return value = ({modpath}, {tmpdir})') 

211 return modpath, tmpdir 

212 

213 

214@begin_end(begin=5, end=5) 

215def python_version() -> t.Tuple: 

216 """ 

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

218 """ 

219 _ = sys.version_info 

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

221 debug.lv5(f'Return value = {retval}') 

222 return retval 

223 

224 

225@begin_end() 

226def raise_on_none(**kwargs): 

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

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

229 if value is None: 

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

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

232 

233 

234@begin_end() 

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

236 """Generate a random string""" 

237 letters = string.ascii_uppercase 

238 if lowercase: 

239 letters = string.ascii_lowercase 

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

241 debug.lv5(f'Return value = {retval}') 

242 return retval 

243 

244 

245@begin_end() 

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

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

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

249 debug.lv5(f'Return value = {retval}') 

250 return retval