Coverage for /Users/buh/.pyenv/versions/3.12.2/envs/pii/lib/python3.12/site-packages/es_pii_tool/helpers/steps.py: 79%

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

2

3 from os import getenv

4 import typing as t

5 import logging

6 from dotmap import DotMap # type: ignore

7 from es_wait import IlmPhase, IlmStep

8 from es_pii_tool.defaults import (

9 PAUSE_DEFAULT,

10 PAUSE_ENVVAR,

11 TIMEOUT_DEFAULT,

12 TIMEOUT_ENVVAR,

13)

14 from es_pii_tool.exceptions import (

15 BadClientResult,

16 FatalError,

17 MissingArgument,

18 MissingError,

19 MissingIndex,

20 ValueMismatch,

21)

22 from es_pii_tool.helpers import elastic_api as api

23 from es_pii_tool.helpers.utils import (

24 configure_ilm_policy,

25 get_alias_actions,

26 strip_ilm_name,

27 es_waiter,

28)

29

30 if t.TYPE_CHECKING:

31 from es_pii_tool.task import Task

32

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

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

35

36 logger = logging.getLogger(__name__)

37

38

39 def 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

51 def 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

72 def 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

86 def 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

95 def 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

127 def 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

143 def 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

164 def 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

180 def 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

210 def 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

230 def 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

254 def 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

312 def 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

318 def 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

331 def 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

337 def 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

343 def 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

356 def 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

370 def 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

378 def 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

394 def 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

448 def 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

457 def 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

488 def 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

494 def 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

504 def 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

510 def 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

529 def 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

535 def 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

548 def 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

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