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

1"""Each function is a single step in PII redaction""" 

2 

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) 

29 

30if t.TYPE_CHECKING: 

31 from es_pii_tool.task import Task 

32 

33PAUSE_VALUE = float(getenv(PAUSE_ENVVAR, default=PAUSE_DEFAULT)) 

34TIMEOUT_VALUE = float(getenv(TIMEOUT_ENVVAR, default=TIMEOUT_DEFAULT)) 

35 

36logger = logging.getLogger(__name__) 

37 

38 

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) 

49 

50 

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) 

70 

71 

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') 

84 

85 

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) 

93 

94 

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.') 

125 

126 

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') 

141 

142 

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') 

162 

163 

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 ) 

178 

179 

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') 

208 

209 

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') 

228 

229 

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) 

248 

249 else: 

250 logger.debug('%s: Index %s has no ILM lifecycle data', stepname, var.index) 

251 log_step(task, stepname, 'end') 

252 

253 

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() 

269 

270 # From here, we check for matching named cloned policy 

271 

272 configure_ilm_policy(task, data) 

273 

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') 

310 

311 

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) 

316 

317 

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 ) 

329 

330 

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) 

335 

336 

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) 

341 

342 

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 ) 

354 

355 

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 ) 

368 

369 

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) 

376 

377 

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) 

392 

393 

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) 

409 

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') 

446 

447 

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) 

455 

456 

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) 

486 

487 

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) 

492 

493 

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 

497 

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) 

502 

503 

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) 

508 

509 

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) 

527 

528 

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) 

533 

534 

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) 

546 

547 

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) 

560 

561 

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')