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

234 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-08-30 20:56 -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 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.Union[str, None] = None) -> None: 

93 """Change/Modify/Update a datastream""" 

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 datastreams. {prettystr(err)}' 

99 ) from err 

100 

101 

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

103 """Create a datastream""" 

104 try: 

105 client.indices.create_data_stream(name=name) 

106 test = Exists(client, name=name, kind='data_stream', pause=PAUSE_VALUE) 

107 test.wait() 

108 except Exception as err: 

109 raise TestbedFailure( 

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

111 ) from err 

112 

113 

114def create_index( 

115 client: 'Elasticsearch', 

116 name: str, 

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

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

119 tier: str = 'hot', 

120) -> None: 

121 """Create named index""" 

122 if not settings: 

123 settings = get_routing(tier=tier) 

124 else: 

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

126 client.indices.create( 

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

128 ) 

129 try: 

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

131 test.wait() 

132 except TimeoutException as err: 

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

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

135 

136 

137def verify( 

138 client: 'Elasticsearch', 

139 kind: str, 

140 name: str, 

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

142) -> bool: 

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

144 success = True 

145 items = ','.split(name) 

146 for item in items: 

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

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

149 success = False 

150 return success 

151 

152 

153def delete( 

154 client: 'Elasticsearch', 

155 kind: str, 

156 name: str, 

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

158) -> bool: 

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

160 which = emap(kind, client) 

161 func = which['delete'] 

162 success = False 

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

164 try: 

165 if kind == 'snapshot': 

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

167 elif kind == 'index': 

168 res = func(index=name) 

169 else: 

170 res = func(name=name) 

171 except NotFoundError as err: 

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

173 success = True 

174 except Exception as err: 

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

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

177 success = True 

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

179 else: 

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

181 else: 

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

183 return success 

184 

185 

186def do_snap( 

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

188) -> None: 

189 """Perform a snapshot""" 

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

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

192 test.wait() 

193 

194 # Mount the index accordingly 

195 client.searchable_snapshots.mount( 

196 repository=repo, 

197 snapshot=snap, 

198 index=idx, 

199 index_settings=get_routing(tier=tier), 

200 renamed_index=mounted_name(idx, tier), 

201 storage=storage_type(tier), 

202 wait_for_completion=True, 

203 ) 

204 # Fix aliases 

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

206 

207 

208def exists( 

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

210) -> bool: 

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

212 if name is None: 

213 return False 

214 retval = True 

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

216 try: 

217 if kind == 'snapshot': 

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

219 elif kind == 'ilm': 

220 retval = func(name=name) 

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

222 retval = func(index=name) 

223 else: 

224 retval = func(name=name) 

225 except NotFoundError: 

226 retval = False 

227 except Exception as err: 

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

229 return retval 

230 

231 

232def fill_index( 

233 client: 'Elasticsearch', 

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

235 doc_generator: t.Union[t.Generator[t.Dict, None, None], None] = None, 

236 options: t.Union[t.Dict, None] = None, 

237) -> None: 

238 """ 

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

240 

241 :param client: ES client 

242 :param name: Index name 

243 :param doc_generator: The generator function 

244 

245 :returns: No return value 

246 """ 

247 if not options: 

248 options = {} 

249 for doc in doc_generator(**options): 

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

251 client.indices.flush(index=name) 

252 client.indices.refresh(index=name) 

253 

254 

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

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

257 retval = None 

258 for alias in get_aliases(client, name): 

259 retval = get_write_index(client, alias) 

260 if retval: 

261 break 

262 return retval 

263 

264 

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

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

267 # Delete the original index 

268 client.indices.delete(index=oldidx) 

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

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

271 

272 

273def get( 

274 client: 'Elasticsearch', 

275 kind: str, 

276 pattern: str, 

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

278) -> t.Sequence[str]: 

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

280 if pattern is None: 

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

282 logger.error(msg) 

283 raise TestbedMisconfig(msg) 

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

285 func = which['get'] 

286 kwargs = which['kwargs'] 

287 if kind == 'snapshot': 

288 kwargs['repository'] = repository 

289 try: 

290 result = func(**kwargs) 

291 except NotFoundError: 

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

293 return [] 

294 except Exception as err: 

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

296 if kind == 'snapshot': 

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

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

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

300 else: 

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

302 retval = list(result.keys()) 

303 return retval 

304 

305 

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

307 """Get aliases from index 'name'""" 

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

309 try: 

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

311 except KeyError: 

312 retval = None 

313 return retval 

314 

315 

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

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

318 resp = resolver(client, name) 

319 data_streams = resp['data_streams'] 

320 retval = [] 

321 if data_streams: 

322 if len(data_streams) > 1: 

323 raise ResultNotExpected( 

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

325 ) 

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

327 return retval 

328 

329 

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

331 """ 

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

333 This is best accomplished by grabbing the last backing_index 

334 """ 

335 backers = get_backing_indices(client, name) 

336 retval = None 

337 if backers: 

338 retval = backers[-1] 

339 return retval 

340 

341 

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

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

344 try: 

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

346 except Exception as err: 

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

348 logger.critical(msg) 

349 raise ResultNotExpected(msg) from err 

350 

351 

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

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

354 ilm = get_ilm(client, name) 

355 try: 

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

357 except KeyError as err: 

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

359 logger.critical(msg) 

360 raise ResultNotExpected(msg) from err 

361 

362 

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

364 """ 

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

366 

367 :param client: A client connection object 

368 :param name: An alias name 

369 

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

371 

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

373 ``is_write_index`` 

374 """ 

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

376 retval = None 

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

378 try: 

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

380 retval = index 

381 break 

382 except KeyError: 

383 continue 

384 return retval 

385 

386 

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

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

389 try: 

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

391 except KeyError: 

392 logger.debug('Index name changed') 

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

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

395 except NotFoundError as err: 

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

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

398 except Exception as err: 

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

400 logger.critical(msg) 

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

402 return retval 

403 

404 

405def ilm_move( 

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

407) -> None: 

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

409 try: 

410 client.ilm.move_to_step( 

411 index=name, current_step=current_step, next_step=next_step 

412 ) 

413 except Exception as err: 

414 msg = ( 

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

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

417 ) 

418 logger.critical(msg) 

419 raise ResultNotExpected(msg, err) 

420 

421 

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

423 """Publish a component template""" 

424 try: 

425 client.cluster.put_component_template( 

426 name=name, template=component, create=True 

427 ) 

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

429 test.wait() 

430 except Exception as err: 

431 raise TestbedFailure( 

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

433 ) from err 

434 

435 

436def put_idx_tmpl( 

437 client, 

438 name: str, 

439 index_patterns: list, 

440 components: list, 

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

442) -> None: 

443 """Publish an index template""" 

444 try: 

445 client.indices.put_index_template( 

446 name=name, 

447 composed_of=components, 

448 data_stream=data_stream, 

449 index_patterns=index_patterns, 

450 create=True, 

451 ) 

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

453 test.wait() 

454 except Exception as err: 

455 raise TestbedFailure( 

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

457 ) from err 

458 

459 

460def put_ilm( 

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

462) -> None: 

463 """Publish an ILM Policy""" 

464 try: 

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

466 except Exception as err: 

467 raise TestbedFailure( 

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

469 ) from err 

470 

471 

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

473 """ 

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

475 

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

477 as an array: 

478 

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

480 

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

482 list 

483 """ 

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

485 

486 

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

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

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

490 

491 

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

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

494 res = {} 

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

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

497 try: 

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

499 except KeyError: 

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

501 retval = None 

502 return retval