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
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-31 13:12 -0600
1"""Utility helper functions"""
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
18logger = logging.getLogger(__name__)
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
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}
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
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
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
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.
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).
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
166 return f"\n{pformat(*args, **kw)}" # newline in front so it always starts clean
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
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
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')
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
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