Coverage for /Users/buh/.pyenv/versions/3.12.2/envs/pii/lib/python3.12/site-packages/es_pii_tool/helpers/steps.py: 79%
313 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-10-01 16:39 -0600
« prev ^ index » next coverage.py v7.5.0, created at 2024-10-01 16:39 -0600
1"""Each function is a single step in PII redaction"""
3from os import getenv
4import typing as t
5import logging
6from dotmap import DotMap # type: ignore
7from es_wait import IlmPhase, IlmStep
8from es_pii_tool.defaults import (
9 PAUSE_DEFAULT,
10 PAUSE_ENVVAR,
11 TIMEOUT_DEFAULT,
12 TIMEOUT_ENVVAR,
13)
14from es_pii_tool.exceptions import (
15 BadClientResult,
16 FatalError,
17 MissingArgument,
18 MissingError,
19 MissingIndex,
20 ValueMismatch,
21)
22from es_pii_tool.helpers import elastic_api as api
23from es_pii_tool.helpers.utils import (
24 configure_ilm_policy,
25 get_alias_actions,
26 strip_ilm_name,
27 es_waiter,
28)
30if t.TYPE_CHECKING:
31 from es_pii_tool.task import Task
33PAUSE_VALUE = float(getenv(PAUSE_ENVVAR, default=PAUSE_DEFAULT))
34TIMEOUT_VALUE = float(getenv(TIMEOUT_ENVVAR, default=TIMEOUT_DEFAULT))
36logger = logging.getLogger(__name__)
39def log_step(task, stepname: str, kind: str):
40 """Function to avoid repetition of code"""
41 msgmap = {
42 'start': 'starting...',
43 'end': 'completed.',
44 'dry-run': 'DRY-RUN. No change will take place',
45 }
46 msg = f'{stepname} {msgmap[kind]}'
47 logger.info(msg)
48 task.add_log(msg)
51def failed_step(task: 'Task', stepname: str, exc):
52 """Function to avoid repetition of code if a step fails"""
53 # MissingIndex, BadClientResult are the only ones inbound
54 upstream = (
55 f'The upstream exception type was {exc.upstream.__name__}, '
56 f'with error message: {exc.upstream.args[0]}'
57 )
58 if isinstance(exc, MissingIndex):
59 msg = f'Step failed because index {exc.missing} was not found. {upstream}'
60 elif isinstance(exc, BadClientResult):
61 msg = (
62 f'Step failed because of a bad or unexpected response or result from '
63 f'the Elasticsearch cluster. {upstream}'
64 )
65 else:
66 msg = f'Step failed for an unexpected reason: {exc}'
67 logger.critical(msg)
68 task.end(False, errors=True, logmsg=f'Failed {stepname}: {msg}')
69 raise FatalError(msg, exc)
72def metastep(task: 'Task', stepname: str, func, *args, **kwargs) -> None:
73 """The reusable step"""
74 log_step(task, stepname, 'start')
75 if not task.job.dry_run:
76 try:
77 func(*args, **kwargs)
78 except (MissingIndex, BadClientResult) as exc:
79 failed_step(task, stepname, exc)
80 else:
81 logger.debug('%s: Dry-Run: No action taken', stepname)
82 log_step(task, stepname, 'dry-run')
83 log_step(task, stepname, 'end')
86def missing_data(stepname, kwargs) -> None:
87 """Avoid duplicated code for data check"""
88 if 'data' not in kwargs:
89 msg = f'"{stepname}" is missing keyword argument(s)'
90 what = 'type: DotMap'
91 names = ['data']
92 raise MissingArgument(msg, what, names)
95def fmwrapper(task: 'Task', stepname: str, var: DotMap) -> None:
96 """Do some task logging around the forcemerge api call"""
97 index = var.redaction_target
98 msg = f'{stepname} Before forcemerge, {api.report_segment_count(var.client, index)}'
99 logger.info(msg)
100 task.add_log(msg)
101 fmkwargs = {}
102 if 'forcemerge' in task.job.config:
103 fmkwargs = task.job.config['forcemerge']
104 fmkwargs['index'] = index
105 if 'only_expunge_deletes' in fmkwargs and fmkwargs['only_expunge_deletes']:
106 msg = 'Forcemerge will only expunge deleted docs!'
107 logger.info(msg)
108 task.add_log(msg)
109 else:
110 mns = 1 # default value
111 if 'max_num_segments' in fmkwargs and isinstance(
112 fmkwargs['max_num_segments'], int
113 ):
114 mns = fmkwargs['max_num_segments']
115 msg = f'Proceeding to forcemerge to {mns} segments per shard'
116 logger.info(msg)
117 task.add_log(msg)
118 logger.debug('forcemerge kwargs = %s', fmkwargs)
119 # Do the actual forcemerging
120 api.forcemerge_index(var.client, **fmkwargs)
121 msg = f'After forcemerge, {api.report_segment_count(var.client, index)}'
122 logger.info(msg)
123 task.add_log(msg)
124 logger.info('Forcemerge completed.')
127def resolve_index(task: 'Task', stepname: str, var: DotMap, **kwargs) -> None:
128 """
129 Resolve the index to see if it's part of a data stream
130 """
131 missing_data(stepname, kwargs)
132 data = kwargs['data']
133 log_step(task, stepname, 'start')
134 result = api.resolve_index(var.client, var.index)
135 logger.debug('resolve data: %s', result)
136 try:
137 data.data_stream = result['indices'][0]['data_stream']
138 except KeyError:
139 logger.debug('%s: Index %s is not part of a data_stream', stepname, var.index)
140 log_step(task, stepname, 'end')
143def pre_delete(task: 'Task', stepname: str, var: DotMap, **kwargs) -> None:
144 """
145 Pre-delete the redacted index to ensure no collisions. Ignore if not present
146 """
147 missing_data(stepname, kwargs)
148 log_step(task, stepname, 'start')
149 if not task.job.dry_run:
150 try:
151 api.delete_index(var.client, var.redaction_target)
152 except MissingIndex:
153 logger.debug(
154 '%s: Pre-delete did not find index "%s"',
155 stepname,
156 var.redaction_target,
157 )
158 # No problem. This is expected.
159 else:
160 log_step(task, stepname, 'dry-run')
161 log_step(task, stepname, 'end')
164def restore_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
165 """Restore index from snapshot"""
166 missing_data(stepname, kwargs)
167 metastep(
168 task,
169 stepname,
170 api.restore_index,
171 var.client,
172 var.repository,
173 var.ss_snap,
174 var.ss_idx,
175 var.redaction_target,
176 index_settings=var.restore_settings.toDict(),
177 )
180def get_index_lifecycle_data(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
181 """
182 Populate data.index with index settings results referenced at
183 INDEXNAME.settings.index.lifecycle
184 """
185 missing_data(stepname, kwargs)
186 data = kwargs['data']
187 log_step(task, stepname, 'start')
188 data.index = DotMap()
189 res = api.get_settings(var.client, var.index)
190 # Set a default value in case we are dealing with non-ILM indices
191 data.index.lifecycle = DotMap(
192 {'name': None, 'rollover_alias': None, 'indexing_complete': True}
193 )
194 try:
195 data.index.lifecycle = DotMap(res[var.index]['settings']['index']['lifecycle'])
196 except KeyError as err:
197 logger.debug(
198 '%s: Index %s missing one or more lifecycle keys: %s',
199 stepname,
200 var.index,
201 err,
202 )
203 if data.index.lifecycle.name:
204 logger.debug('%s: Index lifecycle settings: %s', stepname, data.index.lifecycle)
205 else:
206 logger.debug('%s: Index %s has no ILM lifecycle', stepname, var.index)
207 log_step(task, stepname, 'end')
210def get_ilm_explain_data(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
211 """
212 Populate data.ilm.explain with ilm_explain data
213 """
214 missing_data(stepname, kwargs)
215 data = kwargs['data']
216 log_step(task, stepname, 'start')
217 if data.index.lifecycle.name:
218 data.ilm = DotMap()
219 try:
220 res = api.get_ilm(var.client, var.index)
221 data.ilm.explain = DotMap(res['indices'][var.index])
222 logger.debug('%s: ILM explain settings: %s', stepname, data.ilm.explain)
223 except MissingIndex as exc:
224 failed_step(task, stepname, exc)
225 else:
226 logger.debug('%s: Index %s has no ILM explain data', stepname, var.index)
227 log_step(task, stepname, 'end')
230def get_ilm_lifecycle_data(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
231 """
232 Populate data.ilm.explain with ilm_explain data
233 """
234 missing_data(stepname, kwargs)
235 data = kwargs['data']
236 log_step(task, stepname, 'start')
237 if data.index.lifecycle.name:
238 res = api.get_ilm_lifecycle(var.client, data.index.lifecycle.name)
239 if not res:
240 msg = f'No such ILM policy: {data.index.lifecycle.name}'
241 failed_step(
242 task,
243 stepname,
244 BadClientResult(msg, Exception()),
245 )
246 data.ilm.lifecycle = DotMap(res[data.index.lifecycle.name])
247 logger.debug('%s: ILM lifecycle settings: %s', stepname, data.ilm.lifecycle)
249 else:
250 logger.debug('%s: Index %s has no ILM lifecycle data', stepname, var.index)
251 log_step(task, stepname, 'end')
254def clone_ilm_policy(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
255 """
256 If this index has an ILM policy, we need to clone it so we can attach
257 the new index to it.
258 """
259 missing_data(stepname, kwargs)
260 data = kwargs['data']
261 log_step(task, stepname, 'start')
262 if data.index.lifecycle.name is None or not data.ilm.lifecycle.policy:
263 logger.debug(
264 '%s: Index %s has no ILM lifecycle or policy data', stepname, var.index
265 )
266 log_step(task, stepname, 'end')
267 return
268 data.new = DotMap()
270 # From here, we check for matching named cloned policy
272 configure_ilm_policy(task, data)
274 # New ILM policy naming: pii-tool-POLICYNAME---v###
275 stub = f'pii-tool-{strip_ilm_name(data.index.lifecycle.name)}'
276 policy = data.new.ilmpolicy.toDict() # For comparison
277 resp = {'dummy': 'startval'} # So the while loop can start with something
278 policyver = 0 # Our version number starting point.
279 policymatch = False
280 while resp:
281 data.new.ilmname = f'{stub}---v{policyver + 1:03}'
282 resp = api.get_ilm_lifecycle(var.client, data.new.ilmname) # type: ignore
283 if resp: # We have data, so the name matches
284 # Compare the new policy to the one just returned
285 if policy == resp[data.new.ilmname]['policy']: # type: ignore
286 logger.debug('New policy data matches: %s', data.new.ilmname)
287 policymatch = True
288 break # We can drop out of the loop here.
289 # Implied else: resp has no value, so the while loop will end.
290 policyver += 1
291 logger.debug('New ILM policy name (may already exist): %s', data.new.ilmname)
292 if not task.job.dry_run: # Don't create if dry_run
293 if not policymatch:
294 # Create the cloned ILM policy
295 try:
296 gkw = {'name': data.new.ilmname, 'policy': policy}
297 api.generic_get(var.client.ilm.put_lifecycle, **gkw)
298 except (MissingError, BadClientResult) as exc:
299 logger.error('Unable to put new ILM policy: %s', exc)
300 failed_step(task, stepname, exc)
301 # Implied else: We've arrived at the expected new ILM name
302 # and it does match an existing policy in name and content
303 # so we don't need to create a new one.
304 else:
305 logger.debug(
306 '%s: Dry-Run: ILM policy not created: %s', stepname, data.new.ilmname
307 )
308 log_step(task, stepname, 'dry-run')
309 log_step(task, stepname, 'end')
312def un_ilm_the_restored_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
313 """Remove the lifecycle data from the settings of the restored index"""
314 missing_data(stepname, kwargs)
315 metastep(task, stepname, api.remove_ilm_policy, var.client, var.redaction_target)
318def redact_from_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
319 """Run update by query on new restored index"""
320 missing_data(stepname, kwargs)
321 metastep(
322 task,
323 stepname,
324 api.redact_from_index,
325 var.client,
326 var.redaction_target,
327 task.job.config,
328 )
331def forcemerge_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
332 """Force merge redacted index"""
333 missing_data(stepname, kwargs)
334 metastep(task, stepname, fmwrapper, task, stepname, var)
337def clear_cache(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
338 """Clear cache of redacted index"""
339 missing_data(stepname, kwargs)
340 metastep(task, stepname, api.clear_cache, var.client, var.redaction_target)
343def confirm_redaction(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
344 """Check update by query did its job"""
345 missing_data(stepname, kwargs)
346 metastep(
347 task,
348 stepname,
349 api.check_index,
350 var.client,
351 var.redaction_target,
352 task.job.config,
353 )
356def snapshot_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
357 """Create a new snapshot for mounting our redacted index"""
358 missing_data(stepname, kwargs)
359 metastep(
360 task,
361 stepname,
362 api.take_snapshot,
363 var.client,
364 var.repository,
365 var.new_snap_name,
366 var.redaction_target,
367 )
370def mount_snapshot(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
371 """
372 Mount the index as a searchable snapshot to make the redacted index available
373 """
374 missing_data(stepname, kwargs)
375 metastep(task, stepname, api.mount_index, var)
378def apply_ilm_policy(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
379 """
380 If the index was associated with an ILM policy, associate it with the
381 new, cloned ILM policy.
382 """
383 missing_data(stepname, kwargs)
384 data = kwargs['data']
385 if data.new.ilmname:
386 settings = {'index': {}} # type: ignore
387 # Add all of the original lifecycle settings
388 settings['index']['lifecycle'] = data.index.lifecycle.toDict()
389 # Replace the name with the new ILM policy name
390 settings['index']['lifecycle']['name'] = data.new.ilmname
391 metastep(task, stepname, api.put_settings, var.client, var.mount_name, settings)
394def confirm_ilm_phase(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
395 """
396 Confirm the mounted index is in the expected ILM phase
397 This is done by using move_to_step. If it's already in the step, no problem.
398 If it's in step ``new``, this will advance the index to the expected step.
399 """
400 missing_data(stepname, kwargs)
401 log_step(task, stepname, 'start')
402 # Wait for phase to be "new"
403 waitkw = {'pause': PAUSE_VALUE, 'timeout': TIMEOUT_VALUE}
404 try:
405 es_waiter(var.client, IlmPhase, name=var.mount_name, phase='new', **waitkw)
406 es_waiter(var.client, IlmStep, name=var.mount_name, **waitkw)
407 except BadClientResult as exc:
408 failed_step(task, stepname, exc)
410 try:
411 _ = api.generic_get(var.client.ilm.explain_lifecycle, index=var.mount_name)
412 except MissingError as exc:
413 logger.error('Cannot confirm %s is in phase %s', var.mount_name, var.phase)
414 failed_step(task, stepname, exc)
415 expl = _['indices'][var.mount_name]
416 if not expl['managed']:
417 msg = f'Index {var.mount_name} is not managed by ILM'
418 raise ValueMismatch(msg, expl['managed'], '{"managed": True}')
419 currstep = {'phase': expl['phase'], 'action': expl['action'], 'name': expl['step']}
420 nextstep = {'phase': var.phase, 'action': 'complete', 'name': 'complete'}
421 if not task.job.dry_run: # Don't actually move_to_step if dry_run
422 logger.debug('currstep: %s', currstep)
423 logger.debug('nextstep: %s', nextstep)
424 logger.debug('PHASE: %s', var.phase)
425 try:
426 api.ilm_move(var.client, var.mount_name, currstep, nextstep)
427 except BadClientResult as exc:
428 failed_step(task, stepname, exc)
429 try:
430 es_waiter(
431 var.client, IlmPhase, name=var.mount_name, phase=var.phase, **waitkw
432 )
433 es_waiter(var.client, IlmStep, name=var.mount_name, **waitkw)
434 except BadClientResult as phase_err:
435 msg = f'Unable to wait for ILM step to complete: ERROR :{phase_err}'
436 logger.error(msg)
437 failed_step(task, stepname, phase_err)
438 else:
439 msg = (
440 f'{stepname}: Dry-Run: {var.mount_name} not moved/confirmed to ILM '
441 f'phase {var.phase}'
442 )
443 logger.debug(msg)
444 log_step(task, stepname, 'dry-run')
445 log_step(task, stepname, 'end')
448def delete_redaction_target(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
449 """
450 Now that it's mounted (with a new name), we should delete the redaction_target
451 index
452 """
453 missing_data(stepname, kwargs)
454 metastep(task, stepname, api.delete_index, var.client, var.redaction_target)
457def fixalias_builder(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
458 """This is makes the real fixalias a one liner"""
459 data = kwargs['data']
460 if data.data_stream:
461 msg = f'{stepname} Cannot apply aliases to indices in data_stream'
462 logger.debug(msg)
463 task.add_log(msg)
464 return
465 alias_names = var.aliases.toDict().keys()
466 if not alias_names:
467 msg = f'{stepname} No aliases associated with index {var.index}'
468 task.add_log(msg)
469 logger.warning(msg)
470 else:
471 msg = f'{stepname} Transferring aliases to new index ' f'{var.mount_name}'
472 task.add_log(msg)
473 logger.debug(msg)
474 var.client.indices.update_aliases(
475 actions=get_alias_actions(var.index, var.mount_name, var.aliases.toDict())
476 )
477 verify = var.client.indices.get(index=var.mount_name)[var.mount_name][
478 'aliases'
479 ].keys()
480 if alias_names != verify:
481 msg = f'Alias names do not match! {alias_names} does not match: {verify}'
482 msg2 = f'Failed {stepname}: {msg}'
483 logger.critical(msg2)
484 task.add_log(msg2)
485 raise ValueMismatch(msg, 'alias names mismatch', alias_names)
488def fix_aliases(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
489 """Using the aliases collected from var.index, update mount_name and verify"""
490 missing_data(stepname, kwargs)
491 metastep(task, stepname, fixalias_builder, task, stepname, var, **kwargs)
494def un_ilm_the_original_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
495 """
496 Remove the lifecycle data from the settings of the original index
498 This is chiefly done as a safety measure.
499 """
500 missing_data(stepname, kwargs)
501 metastep(task, stepname, api.remove_ilm_policy, var.client, var.index)
504def close_old_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
505 """Close old mounted snapshot"""
506 missing_data(stepname, kwargs)
507 metastep(task, stepname, api.close_index, var.client, var.index)
510def delete_old_index_builder(task: 'Task', stepname, var: DotMap) -> None:
511 """This makes delete_old_index work with metastep"""
512 if task.job.config['delete']:
513 msg = f'Deleting original mounted index: {var.index}'
514 task.add_log(msg)
515 logger.info(msg)
516 try:
517 api.delete_index(var.client, var.index)
518 except MissingIndex as exc:
519 failed_step(task, stepname, exc)
520 else:
521 msg = (
522 f'delete set to False — not deleting original mounted index: '
523 f'{var.index}'
524 )
525 task.add_log(msg)
526 logger.warning(msg)
529def delete_old_index(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
530 """Delete old mounted snapshot, if configured to do so"""
531 missing_data(stepname, kwargs)
532 metastep(task, stepname, delete_old_index_builder, task, stepname, var)
535def assign_aliases(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
536 """Put the starting index name on new mounted index as alias"""
537 missing_data(stepname, kwargs)
538 data = kwargs['data']
539 if data.data_stream:
540 log_step(task, stepname, 'start')
541 msg = f'{stepname}: Cannot apply aliases to indices in data_stream'
542 logger.debug(msg)
543 log_step(task, stepname, 'end')
544 return
545 metastep(task, stepname, api.assign_alias, var.client, var.mount_name, var.index)
548def reassociate_index_with_ds(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
549 """
550 If the index was associated with a data_stream, reassociate it with the
551 data_stream again.
552 """
553 missing_data(stepname, kwargs)
554 data = kwargs['data']
555 acts = [{'add_backing_index': {'index': var.mount_name}}]
556 if data.data_stream:
557 acts[0]['add_backing_index']['data_stream'] = data.data_stream
558 logger.debug('%s: Modify data_stream actions: %s', stepname, acts)
559 metastep(task, stepname, api.modify_data_stream, var.client, acts)
562def record_it(task: 'Task', stepname, var: DotMap, **kwargs) -> None:
563 """Record the now-deletable snapshot in the job's tracking index."""
564 missing_data(stepname, kwargs)
565 log_step(task, stepname, 'start')
566 task.job.cleanup.append(var.ss_snap)
567 log_step(task, stepname, 'end')