Coverage for /Users/buh/.pyenv/versions/3.12.9/envs/es-testbed/lib/python3.12/site-packages/es_testbed/helpers/es_api.py: 99%

241 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-03-17 19:30 -0600

1"""Functions that make Elasticsearch API Calls""" 

2 

3import typing as t 

4import logging 

5from os import getenv 

6from elasticsearch8.exceptions import NotFoundError, TransportError 

7from es_wait import Exists, Snapshot 

8from es_wait.exceptions import EsWaitFatal, EsWaitTimeout 

9from ..defaults import MAPPING, PAUSE_DEFAULT, PAUSE_ENVVAR 

10from ..exceptions import ( 

11 NameChanged, 

12 ResultNotExpected, 

13 TestbedFailure, 

14 TestbedMisconfig, 

15) 

16from ..helpers.utils import ( 

17 get_routing, 

18 mounted_name, 

19 prettystr, 

20 storage_type, 

21) 

22 

23if t.TYPE_CHECKING: 

24 from elasticsearch8 import Elasticsearch 

25 

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

27 

28logger = logging.getLogger(__name__) 

29 

30# pylint: disable=W0707 

31 

32 

33def emap(kind: str, es: 'Elasticsearch', value=None) -> t.Dict[str, t.Any]: 

34 """Return a value from a dictionary""" 

35 _ = { 

36 'alias': { 

37 'delete': es.indices.delete_alias, 

38 'exists': es.indices.exists_alias, 

39 'get': es.indices.get_alias, 

40 'kwargs': {'index': value, 'expand_wildcards': ['open', 'closed']}, 

41 'plural': 'alias(es)', 

42 }, 

43 'data_stream': { 

44 'delete': es.indices.delete_data_stream, 

45 'exists': es.indices.exists, 

46 'get': es.indices.get_data_stream, 

47 'kwargs': {'name': value, 'expand_wildcards': ['open', 'closed']}, 

48 'plural': 'data_stream(s)', 

49 'key': 'data_streams', 

50 }, 

51 'index': { 

52 'delete': es.indices.delete, 

53 'exists': es.indices.exists, 

54 'get': es.indices.get, 

55 'kwargs': {'index': value, 'expand_wildcards': ['open', 'closed']}, 

56 'plural': 'index(es)', 

57 }, 

58 'template': { 

59 'delete': es.indices.delete_index_template, 

60 'exists': es.indices.exists_index_template, 

61 'get': es.indices.get_index_template, 

62 'kwargs': {'name': value}, 

63 'plural': 'index template(s)', 

64 'key': 'index_templates', 

65 }, 

66 'ilm': { 

67 'delete': es.ilm.delete_lifecycle, 

68 'exists': es.ilm.get_lifecycle, 

69 'get': es.ilm.get_lifecycle, 

70 'kwargs': {'name': value}, 

71 'plural': 'ilm policy(ies)', 

72 }, 

73 'component': { 

74 'delete': es.cluster.delete_component_template, 

75 'exists': es.cluster.exists_component_template, 

76 'get': es.cluster.get_component_template, 

77 'kwargs': {'name': value}, 

78 'plural': 'component template(s)', 

79 'key': 'component_templates', 

80 }, 

81 'snapshot': { 

82 'delete': es.snapshot.delete, 

83 'exists': es.snapshot.get, 

84 'get': es.snapshot.get, 

85 'kwargs': {'snapshot': value}, 

86 'plural': 'snapshot(s)', 

87 }, 

88 } 

89 return _[kind] 

90 

91 

92def change_ds(client: 'Elasticsearch', actions: t.Optional[str] = None) -> None: 

93 """Change/Modify/Update a data_stream""" 

94 try: 

95 client.indices.modify_data_stream(actions=actions, body=None) 

96 except Exception as err: 

97 raise ResultNotExpected( 

98 f'Unable to modify data_streams. {prettystr(err)}' 

99 ) from err 

100 

101 

102# pylint: disable=R0913,R0917 

103def wait_wrapper( 

104 client: 'Elasticsearch', 

105 wait_cls: t.Callable, 

106 wait_kwargs: t.Dict, 

107 func: t.Callable, 

108 f_kwargs: t.Dict, 

109) -> None: 

110 """Wrapper function for waiting on an object to be created""" 

111 try: 

112 func(**f_kwargs) 

113 test = wait_cls(client, **wait_kwargs) 

114 test.wait() 

115 except EsWaitFatal as wait: 

116 msg = f'{wait.message}. Elapsed time: {wait.elapsed}. Errors: {wait.errors}' 

117 raise TestbedFailure(msg) from wait 

118 except EsWaitTimeout as wait: 

119 msg = f'{wait.message}. Elapsed time: {wait.elapsed}. Timeout: {wait.timeout}' 

120 raise TestbedFailure(msg) from wait 

121 except TransportError as err: 

122 raise TestbedFailure( 

123 f'Elasticsearch TransportError class exception encountered:' 

124 f'{prettystr(err)}' 

125 ) from err 

126 except Exception as err: 

127 raise TestbedFailure(f'General Exception caught: {prettystr(err)}') from err 

128 

129 

130def create_data_stream(client: 'Elasticsearch', name: str) -> None: 

131 """Create a data_stream""" 

132 wait_kwargs = {'name': name, 'kind': 'data_stream', 'pause': PAUSE_VALUE} 

133 f_kwargs = {'name': name} 

134 wait_wrapper( 

135 client, Exists, wait_kwargs, client.indices.create_data_stream, f_kwargs 

136 ) 

137 

138 

139def create_index( 

140 client: 'Elasticsearch', 

141 name: str, 

142 aliases: t.Union[t.Dict, None] = None, 

143 settings: t.Union[t.Dict, None] = None, 

144 tier: str = 'hot', 

145) -> None: 

146 """Create named index""" 

147 if not settings: 

148 settings = get_routing(tier=tier) 

149 else: 

150 settings.update(get_routing(tier=tier)) 

151 wait_kwargs = {'name': name, 'kind': 'index', 'pause': PAUSE_VALUE} 

152 f_kwargs = { 

153 'index': name, 

154 'aliases': aliases, 

155 'mappings': MAPPING, 

156 'settings': settings, 

157 } 

158 wait_wrapper(client, Exists, wait_kwargs, client.indices.create, f_kwargs) 

159 return exists(client, 'index', name) 

160 

161 

162def verify( 

163 client: 'Elasticsearch', 

164 kind: str, 

165 name: str, 

166 repository: t.Optional[str] = None, 

167) -> bool: 

168 """Verify that whatever was deleted is actually deleted""" 

169 success = True 

170 items = name.split(',') 

171 for item in items: 

172 result = exists(client, kind, item, repository=repository) 

173 if result: # That means it's still in the cluster 

174 success = False 

175 return success 

176 

177 

178def delete( 

179 client: 'Elasticsearch', 

180 kind: str, 

181 name: str, 

182 repository: t.Optional[str] = None, 

183) -> bool: 

184 """Delete the named object of type kind""" 

185 which = emap(kind, client) 

186 func = which['delete'] 

187 success = False 

188 if name is not None: # Typically only with ilm 

189 try: 

190 if kind == 'snapshot': 

191 res = func(snapshot=name, repository=repository) 

192 elif kind == 'index': 

193 res = func(index=name) 

194 else: 

195 res = func(name=name) 

196 except NotFoundError as err: 

197 logger.warning(f'{kind} named {name} not found: {prettystr(err)}') 

198 return True 

199 except Exception as err: 

200 raise ResultNotExpected(f'Unexpected result: {prettystr(err)}') from err 

201 if 'acknowledged' in res and res['acknowledged']: 

202 success = True 

203 logger.info(f'Deleted {which["plural"]}: "{name}"') 

204 else: 

205 success = verify(client, kind, name, repository=repository) 

206 else: 

207 logger.debug(f'"{kind}" has a None value for name') 

208 return success 

209 

210 

211def do_snap( 

212 client: 'Elasticsearch', repo: str, snap: str, idx: str, tier: str = 'cold' 

213) -> None: 

214 """Perform a snapshot""" 

215 wait_kwargs = {'snapshot': snap, 'repository': repo, 'pause': 1, 'timeout': 60} 

216 f_kwargs = {'repository': repo, 'snapshot': snap, 'indices': idx} 

217 wait_wrapper(client, Snapshot, wait_kwargs, client.snapshot.create, f_kwargs) 

218 

219 # Mount the index accordingly 

220 client.searchable_snapshots.mount( 

221 repository=repo, 

222 snapshot=snap, 

223 index=idx, 

224 index_settings=get_routing(tier=tier), 

225 renamed_index=mounted_name(idx, tier), 

226 storage=storage_type(tier), 

227 wait_for_completion=True, 

228 ) 

229 # Fix aliases 

230 fix_aliases(client, idx, mounted_name(idx, tier)) 

231 

232 

233def exists( 

234 client: 'Elasticsearch', kind: str, name: str, repository: t.Union[str, None] = None 

235) -> bool: 

236 """Return boolean existence of the named kind of object""" 

237 if name is None: 

238 return False 

239 retval = True 

240 func = emap(kind, client)['exists'] 

241 try: 

242 if kind == 'snapshot': 

243 # Expected response: {'snapshots': [{'snapshot': name, ...}]} 

244 # Since we are specifying by name, there should only be one returned 

245 res = func(snapshot=name, repository=repository) 

246 logger.debug(f'Snapshot response: {res}') 

247 # If there are no entries, load a default None value for the check 

248 _ = dict(res['snapshots'][0]) if res else {'snapshot': None} 

249 # Since there should be only 1 snapshot with this name, we can check it 

250 retval = bool(_['snapshot'] == name) 

251 elif kind == 'ilm': 

252 # There is no true 'exists' method for ILM, so we have to get the policy 

253 # and check for a NotFoundError 

254 retval = bool(name in dict(func(name=name))) 

255 elif kind in ['index', 'data_stream']: 

256 retval = func(index=name) 

257 else: 

258 retval = func(name=name) 

259 except NotFoundError: 

260 retval = False 

261 except Exception as err: 

262 raise ResultNotExpected(f'Unexpected result: {prettystr(err)}') from err 

263 return retval 

264 

265 

266def fill_index( 

267 client: 'Elasticsearch', 

268 name: t.Optional[str] = None, 

269 doc_generator: t.Optional[t.Generator[t.Dict, None, None]] = None, 

270 options: t.Optional[t.Dict] = None, 

271) -> None: 

272 """ 

273 Create and fill the named index with mappings and settings as directed 

274 

275 :param client: ES client 

276 :param name: Index name 

277 :param doc_generator: The generator function 

278 

279 :returns: No return value 

280 """ 

281 if not options: 

282 options = {} 

283 for doc in doc_generator(**options): 

284 client.index(index=name, document=doc) 

285 client.indices.flush(index=name) 

286 client.indices.refresh(index=name) 

287 

288 

289def find_write_index(client: 'Elasticsearch', name: str) -> t.AnyStr: 

290 """Find the write_index for an alias by searching any index the alias points to""" 

291 retval = None 

292 for alias in get_aliases(client, name): 

293 retval = get_write_index(client, alias) 

294 if retval: 

295 break 

296 return retval 

297 

298 

299def fix_aliases(client: 'Elasticsearch', oldidx: str, newidx: str) -> None: 

300 """Fix aliases using the new and old index names as data""" 

301 # Delete the original index 

302 client.indices.delete(index=oldidx) 

303 # Add the original index name as an alias to the mounted index 

304 client.indices.put_alias(index=f'{newidx}', name=oldidx) 

305 

306 

307def get( 

308 client: 'Elasticsearch', 

309 kind: str, 

310 pattern: str, 

311 repository: t.Optional[str] = None, 

312) -> t.Sequence[str]: 

313 """get any/all objects of type kind matching pattern""" 

314 if pattern is None: 

315 msg = f'"{kind}" has a None value for pattern' 

316 logger.error(msg) 

317 raise TestbedMisconfig(msg) 

318 which = emap(kind, client, value=pattern) 

319 func = which['get'] 

320 kwargs = which['kwargs'] 

321 if kind == 'snapshot': 

322 kwargs['repository'] = repository 

323 try: 

324 result = func(**kwargs) 

325 except NotFoundError: 

326 logger.debug(f'{kind} pattern "{pattern}" had zero matches') 

327 return [] 

328 except Exception as err: 

329 raise ResultNotExpected(f'Unexpected result: {prettystr(err)}') from err 

330 if kind == 'snapshot': 

331 retval = [x['snapshot'] for x in result['snapshots']] 

332 elif kind in ['data_stream', 'template', 'component']: 

333 retval = [x['name'] for x in result[which['key']]] 

334 else: 

335 # ['alias', 'ilm', 'index'] 

336 retval = list(result.keys()) 

337 return retval 

338 

339 

340def get_aliases(client: 'Elasticsearch', name: str) -> t.Sequence[str]: 

341 """Get aliases from index 'name'""" 

342 res = client.indices.get(index=name) 

343 try: 

344 retval = list(res[name]['aliases'].keys()) 

345 except KeyError: 

346 retval = None 

347 return retval 

348 

349 

350def get_backing_indices(client: 'Elasticsearch', name: str) -> t.Sequence[str]: 

351 """Get the backing indices from the named data_stream""" 

352 resp = resolver(client, name) 

353 data_streams = resp['data_streams'] 

354 retval = [] 

355 if data_streams: 

356 if len(data_streams) > 1: 

357 raise ResultNotExpected( 

358 f'Expected only a single data_stream matching {name}' 

359 ) 

360 retval = data_streams[0]['backing_indices'] 

361 return retval 

362 

363 

364def get_ds_current(client: 'Elasticsearch', name: str) -> str: 

365 """ 

366 Find which index is the current 'write' index of the data_stream 

367 This is best accomplished by grabbing the last backing_index 

368 """ 

369 backers = get_backing_indices(client, name) 

370 retval = None 

371 if backers: 

372 retval = backers[-1] 

373 return retval 

374 

375 

376def get_ilm(client: 'Elasticsearch', pattern: str) -> t.Union[t.Dict[str, str], None]: 

377 """Get any ILM entity in ES that matches pattern""" 

378 try: 

379 return client.ilm.get_lifecycle(name=pattern) 

380 except Exception as err: 

381 msg = f'Unable to get ILM lifecycle matching {pattern}. Error: {prettystr(err)}' 

382 logger.critical(msg) 

383 raise ResultNotExpected(msg) from err 

384 

385 

386def get_ilm_phases(client: 'Elasticsearch', name: str) -> t.Dict: 

387 """Return the policy/phases part of the ILM policy identified by 'name'""" 

388 ilm = get_ilm(client, name) 

389 try: 

390 return ilm[name]['policy']['phases'] 

391 except KeyError as err: 

392 msg = f'Unable to get ILM lifecycle named {name}. Error: {prettystr(err)}' 

393 logger.critical(msg) 

394 raise ResultNotExpected(msg) from err 

395 

396 

397def get_write_index(client: 'Elasticsearch', name: str) -> str: 

398 """ 

399 Calls :py:meth:`~.elasticsearch.client.IndicesClient.get_alias` 

400 

401 :param client: A client connection object 

402 :param name: An alias name 

403 

404 :type client: :py:class:`~.elasticsearch.Elasticsearch` 

405 

406 :returns: The the index name associated with the alias that is designated 

407 ``is_write_index`` 

408 """ 

409 response = client.indices.get_alias(index=name) 

410 retval = None 

411 for index in list(response.keys()): 

412 try: 

413 if response[index]['aliases'][name]['is_write_index']: 

414 retval = index 

415 break 

416 except KeyError: 

417 continue 

418 return retval 

419 

420 

421def ilm_explain(client: 'Elasticsearch', name: str) -> t.Union[t.Dict, None]: 

422 """Return the results from the ILM Explain API call for the named index""" 

423 try: 

424 retval = client.ilm.explain_lifecycle(index=name)['indices'][name] 

425 except KeyError: 

426 logger.debug('Index name changed') 

427 new = list(client.ilm.explain_lifecycle(index=name)['indices'].keys())[0] 

428 retval = client.ilm.explain_lifecycle(index=new)['indices'][new] 

429 except NotFoundError as err: 

430 logger.warning(f'Datastream/Index Name changed. {name} was not found') 

431 raise NameChanged(f'{name} was not found, likely due to a name change') from err 

432 except Exception as err: 

433 msg = f'Unable to get ILM information for index {name}' 

434 logger.critical(msg) 

435 raise ResultNotExpected(f'{msg}. Exception: {prettystr(err)}') from err 

436 return retval 

437 

438 

439def ilm_move( 

440 client: 'Elasticsearch', name: str, current_step: t.Dict, next_step: t.Dict 

441) -> None: 

442 """Move index 'name' from the current step to the next step""" 

443 try: 

444 client.ilm.move_to_step( 

445 index=name, current_step=current_step, next_step=next_step 

446 ) 

447 except Exception as err: 

448 msg = ( 

449 f'Unable to move index {name} to ILM next step: {next_step}. ' 

450 f'Error: {prettystr(err)}' 

451 ) 

452 logger.critical(msg) 

453 raise ResultNotExpected(msg, (err,)) 

454 

455 

456def put_comp_tmpl(client: 'Elasticsearch', name: str, component: t.Dict) -> None: 

457 """Publish a component template""" 

458 wait_kwargs = {'name': name, 'kind': 'component_template', 'pause': PAUSE_VALUE} 

459 f_kwargs = {'name': name, 'template': component, 'create': True} 

460 wait_wrapper( 

461 client, 

462 Exists, 

463 wait_kwargs, 

464 client.cluster.put_component_template, 

465 f_kwargs, 

466 ) 

467 

468 

469def put_idx_tmpl( 

470 client: 'Elasticsearch', 

471 name: str, 

472 index_patterns: t.List[str], 

473 components: t.List[str], 

474 data_stream: t.Optional[t.Dict] = None, 

475) -> None: 

476 """Publish an index template""" 

477 wait_kwargs = {'name': name, 'kind': 'index_template', 'pause': PAUSE_VALUE} 

478 f_kwargs = { 

479 'name': name, 

480 'composed_of': components, 

481 'data_stream': data_stream, 

482 'index_patterns': index_patterns, 

483 'create': True, 

484 } 

485 wait_wrapper( 

486 client, 

487 Exists, 

488 wait_kwargs, 

489 client.indices.put_index_template, 

490 f_kwargs, 

491 ) 

492 

493 

494def put_ilm( 

495 client: 'Elasticsearch', name: str, policy: t.Union[t.Dict, None] = None 

496) -> None: 

497 """Publish an ILM Policy""" 

498 try: 

499 client.ilm.put_lifecycle(name=name, policy=policy) 

500 except Exception as err: 

501 raise TestbedFailure( 

502 f'Unable to put index lifecycle policy {name}. Error: {prettystr(err)}' 

503 ) from err 

504 

505 

506def resolver(client: 'Elasticsearch', name: str) -> dict: 

507 """ 

508 Resolve details about the entity, be it an index, alias, or data_stream 

509 

510 Because you can pass search patterns and aliases as name, each element comes back 

511 as an array: 

512 

513 {'indices': [], 'aliases': [], 'data_streams': []} 

514 

515 If you only resolve a single index or data stream, you will still have a 1-element 

516 list 

517 """ 

518 return client.indices.resolve_index(name=name, expand_wildcards=['open', 'closed']) 

519 

520 

521def rollover(client: 'Elasticsearch', name: str) -> None: 

522 """Rollover alias or data_stream identified by name""" 

523 client.indices.rollover(alias=name, wait_for_active_shards='all') 

524 

525 

526def snapshot_name(client: 'Elasticsearch', name: str) -> t.Union[t.AnyStr, None]: 

527 """Get the name of the snapshot behind the mounted index data""" 

528 res = {} 

529 if exists(client, 'index', name): # Can jump straight to nested keys if it exists 

530 res = client.indices.get(index=name)[name]['settings']['index'] 

531 try: 

532 retval = res['store']['snapshot']['snapshot_name'] 

533 except KeyError: 

534 logger.error(f'{name} is not a searchable snapshot') 

535 retval = None 

536 return retval