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

232 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-05-03 08:16 -0600

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

2 

3import typing as t 

4import logging 

5from os import getenv 

6from elasticsearch8.exceptions import NotFoundError 

7from es_wait import Exists, Snapshot 

8from ..defaults import MAPPING, PAUSE_DEFAULT, PAUSE_ENVVAR 

9from ..exceptions import ( 

10 NameChanged, 

11 ResultNotExpected, 

12 TestbedFailure, 

13 TestbedMisconfig, 

14 TimeoutException, 

15) 

16from ..helpers.utils import ( 

17 doc_gen, 

18 get_routing, 

19 mounted_name, 

20 prettystr, 

21 storage_type, 

22) 

23 

24if t.TYPE_CHECKING: 

25 from elasticsearch8 import Elasticsearch 

26 

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

28 

29logger = logging.getLogger(__name__) 

30 

31# pylint: disable=W0707 

32 

33 

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

35 """Return a value from a dictionary""" 

36 _ = { 

37 'alias': { 

38 'delete': es.indices.delete_alias, 

39 'exists': es.indices.exists_alias, 

40 'get': es.indices.get_alias, 

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

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

43 }, 

44 'data_stream': { 

45 'delete': es.indices.delete_data_stream, 

46 'exists': es.indices.exists, 

47 'get': es.indices.get_data_stream, 

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

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

50 'key': 'data_streams', 

51 }, 

52 'index': { 

53 'delete': es.indices.delete, 

54 'exists': es.indices.exists, 

55 'get': es.indices.get, 

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

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

58 }, 

59 'template': { 

60 'delete': es.indices.delete_index_template, 

61 'exists': es.indices.exists_index_template, 

62 'get': es.indices.get_index_template, 

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

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

65 'key': 'index_templates', 

66 }, 

67 'ilm': { 

68 'delete': es.ilm.delete_lifecycle, 

69 'exists': es.ilm.get_lifecycle, 

70 'get': es.ilm.get_lifecycle, 

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

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

73 }, 

74 'component': { 

75 'delete': es.cluster.delete_component_template, 

76 'exists': es.cluster.exists_component_template, 

77 'get': es.cluster.get_component_template, 

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

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

80 'key': 'component_templates', 

81 }, 

82 'snapshot': { 

83 'delete': es.snapshot.delete, 

84 'exists': es.snapshot.get, 

85 'get': es.snapshot.get, 

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

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

88 }, 

89 } 

90 return _[kind] 

91 

92 

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

94 """Change/Modify/Update a datastream""" 

95 try: 

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

97 except Exception as err: 

98 raise ResultNotExpected( 

99 f'Unable to modify datastreams. {prettystr(err)}' 

100 ) from err 

101 

102 

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

104 """Create a datastream""" 

105 try: 

106 client.indices.create_data_stream(name=name) 

107 test = Exists(client, name=name, kind='datastream', pause=PAUSE_VALUE) 

108 test.wait_for_it() 

109 except Exception as err: 

110 raise TestbedFailure( 

111 f'Unable to create datastream {name}. Error: {prettystr(err)}' 

112 ) from err 

113 

114 

115def create_index( 

116 client: 'Elasticsearch', 

117 name: str, 

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

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

120 tier: str = 'hot', 

121) -> None: 

122 """Create named index""" 

123 if not settings: 

124 settings = get_routing(tier=tier) 

125 else: 

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

127 client.indices.create( 

128 index=name, aliases=aliases, mappings=MAPPING, settings=settings 

129 ) 

130 try: 

131 test = Exists(client, name=name, kind='index', pause=PAUSE_VALUE) 

132 test.wait_for_it() 

133 except TimeoutException as err: 

134 raise ResultNotExpected(f'Failed to create index {name}') from err 

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

136 

137 

138def verify( 

139 client: 'Elasticsearch', 

140 kind: str, 

141 name: str, 

142 repository: t.Union[str, None] = None, 

143) -> bool: 

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

145 success = True 

146 items = ','.split(name) 

147 for item in items: 

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

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

150 success = False 

151 return success 

152 

153 

154def delete( 

155 client: 'Elasticsearch', 

156 kind: str, 

157 name: str, 

158 repository: t.Union[str, None] = None, 

159) -> bool: 

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

161 which = emap(kind, client) 

162 func = which['delete'] 

163 success = False 

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

165 try: 

166 if kind == 'snapshot': 

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

168 elif kind == 'index': 

169 res = func(index=name) 

170 else: 

171 res = func(name=name) 

172 except NotFoundError as err: 

173 logger.warning('%s named %s not found: %s', kind, name, prettystr(err)) 

174 success = True 

175 except Exception as err: 

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

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

178 success = True 

179 logger.info('Deleted %s: "%s"', which['plural'], name) 

180 else: 

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

182 else: 

183 logger.debug('"%s" has a None value for name', kind) 

184 return success 

185 

186 

187def do_snap( 

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

189) -> None: 

190 """Perform a snapshot""" 

191 client.snapshot.create(repository=repo, snapshot=snap, indices=idx) 

192 test = Snapshot(client, snapshot=snap, repository=repo, pause=1, timeout=60) 

193 test.wait_for_it() 

194 

195 # Mount the index accordingly 

196 client.searchable_snapshots.mount( 

197 repository=repo, 

198 snapshot=snap, 

199 index=idx, 

200 index_settings=get_routing(tier=tier), 

201 renamed_index=mounted_name(idx, tier), 

202 storage=storage_type(tier), 

203 wait_for_completion=True, 

204 ) 

205 # Fix aliases 

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

207 

208 

209def exists( 

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

211) -> bool: 

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

213 if name is None: 

214 return False 

215 retval = True 

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

217 try: 

218 if kind == 'snapshot': 

219 retval = func(snapshot=name, repository=repository) 

220 elif kind == 'ilm': 

221 retval = func(name=name) 

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

223 retval = func(index=name) 

224 else: 

225 retval = func(name=name) 

226 except NotFoundError: 

227 retval = False 

228 except Exception as err: 

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

230 return retval 

231 

232 

233def fill_index( 

234 client: 'Elasticsearch', 

235 name: t.Union[str, None] = None, 

236 count: t.Union[int, None] = None, 

237 start_num: t.Union[int, None] = None, 

238 match: bool = True, 

239) -> None: 

240 """ 

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

242 

243 :param client: ES client 

244 :param name: Index name 

245 :param count: The number of docs to create 

246 :param start_number: Where to start the incrementing number 

247 :param match: Whether to use the default values for key (True) or random strings 

248 (False) 

249 

250 :type client: es 

251 :type name: str 

252 :type count: int 

253 :type start_number: int 

254 :type match: bool 

255 

256 :rtype: None 

257 :returns: No return value 

258 """ 

259 for doc in doc_gen(count=count, start_at=start_num, match=match): 

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

261 client.indices.flush(index=name) 

262 client.indices.refresh(index=name) 

263 

264 

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

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

267 retval = None 

268 for alias in get_aliases(client, name): 

269 retval = get_write_index(client, alias) 

270 if retval: 

271 break 

272 return retval 

273 

274 

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

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

277 # Delete the original index 

278 client.indices.delete(index=oldidx) 

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

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

281 

282 

283def get( 

284 client: 'Elasticsearch', 

285 kind: str, 

286 pattern: str, 

287 repository: t.Union[str, None] = None, 

288) -> t.Sequence[str]: 

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

290 if pattern is None: 

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

292 logger.error(msg) 

293 raise TestbedMisconfig(msg) 

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

295 func = which['get'] 

296 kwargs = which['kwargs'] 

297 if kind == 'snapshot': 

298 kwargs['repository'] = repository 

299 try: 

300 result = func(**kwargs) 

301 except NotFoundError: 

302 logger.debug('%s pattern "%s" had zero matches', kind, pattern) 

303 return [] 

304 except Exception as err: 

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

306 if kind == 'snapshot': 

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

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

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

310 else: 

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

312 retval = list(result.keys()) 

313 return retval 

314 

315 

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

317 """Get aliases from index 'name'""" 

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

319 try: 

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

321 except KeyError: 

322 retval = None 

323 return retval 

324 

325 

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

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

328 resp = resolver(client, name) 

329 data_streams = resp['data_streams'] 

330 retval = [] 

331 if data_streams: 

332 if len(data_streams) > 1: 

333 raise ResultNotExpected( 

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

335 ) 

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

337 return retval 

338 

339 

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

341 """ 

342 Find which index is the current 'write' index of the datastream 

343 This is best accomplished by grabbing the last backing_index 

344 """ 

345 backers = get_backing_indices(client, name) 

346 retval = None 

347 if backers: 

348 retval = backers[-1] 

349 return retval 

350 

351 

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

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

354 try: 

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

356 except Exception as err: 

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

358 logger.critical(msg) 

359 raise ResultNotExpected(msg) from err 

360 

361 

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

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

364 ilm = get_ilm(client, name) 

365 try: 

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

367 except KeyError as err: 

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

369 logger.critical(msg) 

370 raise ResultNotExpected(msg) from err 

371 

372 

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

374 """ 

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

376 

377 :param client: A client connection object 

378 :param name: An alias name 

379 

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

381 

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

383 ``is_write_index`` 

384 """ 

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

386 retval = None 

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

388 try: 

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

390 retval = index 

391 break 

392 except KeyError: 

393 continue 

394 return retval 

395 

396 

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

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

399 try: 

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

401 except KeyError: 

402 logger.debug('Index name changed') 

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

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

405 except NotFoundError as err: 

406 logger.warning('Datastream/Index Name changed. %s was not found', name) 

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

408 except Exception as err: 

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

410 logger.critical(msg) 

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

412 return retval 

413 

414 

415def ilm_move( 

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

417) -> None: 

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

419 try: 

420 client.ilm.move_to_step( 

421 index=name, current_step=current_step, next_step=next_step 

422 ) 

423 except Exception as err: 

424 msg = ( 

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

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

427 ) 

428 logger.critical(msg) 

429 raise ResultNotExpected(msg, err) 

430 

431 

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

433 """Publish a component template""" 

434 try: 

435 client.cluster.put_component_template( 

436 name=name, template=component, create=True 

437 ) 

438 test = Exists(client, name=name, kind='component', pause=PAUSE_VALUE) 

439 test.wait_for_it() 

440 except Exception as err: 

441 raise TestbedFailure( 

442 f'Unable to create component template {name}. Error: {prettystr(err)}' 

443 ) from err 

444 

445 

446def put_idx_tmpl( 

447 client, 

448 name: str, 

449 index_patterns: list, 

450 components: list, 

451 data_stream: t.Union[t.Dict, None] = None, 

452) -> None: 

453 """Publish an index template""" 

454 try: 

455 client.indices.put_index_template( 

456 name=name, 

457 composed_of=components, 

458 data_stream=data_stream, 

459 index_patterns=index_patterns, 

460 create=True, 

461 ) 

462 test = Exists(client, name=name, kind='template', pause=PAUSE_VALUE) 

463 test.wait_for_it() 

464 except Exception as err: 

465 raise TestbedFailure( 

466 f'Unable to create index template {name}. Error: {prettystr(err)}' 

467 ) from err 

468 

469 

470def put_ilm( 

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

472) -> None: 

473 """Publish an ILM Policy""" 

474 try: 

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

476 except Exception as err: 

477 raise TestbedFailure( 

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

479 ) from err 

480 

481 

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

483 """ 

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

485 

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

487 as an array: 

488 

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

490 

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

492 list 

493 """ 

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

495 

496 

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

498 """Rollover alias or datastream identified by name""" 

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

500 

501 

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

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

504 res = {} 

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

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

507 try: 

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

509 except KeyError: 

510 logger.error('%s is not a searchable snapshot') 

511 retval = None 

512 return retval