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
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 23:00 -0600
1"""Utility helper functions"""
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
12logger = logging.getLogger(__name__)
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
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}
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()
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 }
95def getlogger(name: str) -> logging.getLogger:
96 """Return a named logger"""
97 return logging.getLogger(name)
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}
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.
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
143def mapping_component() -> t.Dict:
144 """Return a mappings component template"""
145 return {'mappings': MAPPING}
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}'
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.
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
176 return f"\n{pformat(*args, **kw)}" # newline in front so it always starts clean
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)))
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
199def storage_type(tier: str) -> t.Dict:
200 """Return the storage type of a searchable snapshot by tier"""
201 return TIER[tier]["storage"]