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
« prev ^ index » next coverage.py v7.6.12, created at 2025-03-17 19:30 -0600
1"""Functions that make Elasticsearch API Calls"""
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)
23if t.TYPE_CHECKING:
24 from elasticsearch8 import Elasticsearch
26PAUSE_VALUE = float(getenv(PAUSE_ENVVAR, default=PAUSE_DEFAULT))
28logger = logging.getLogger(__name__)
30# pylint: disable=W0707
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]
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
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
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 )
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)
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
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
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)
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))
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
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
275 :param client: ES client
276 :param name: Index name
277 :param doc_generator: The generator function
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)
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
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)
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
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
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
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
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
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
397def get_write_index(client: 'Elasticsearch', name: str) -> str:
398 """
399 Calls :py:meth:`~.elasticsearch.client.IndicesClient.get_alias`
401 :param client: A client connection object
402 :param name: An alias name
404 :type client: :py:class:`~.elasticsearch.Elasticsearch`
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
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
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,))
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 )
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 )
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
506def resolver(client: 'Elasticsearch', name: str) -> dict:
507 """
508 Resolve details about the entity, be it an index, alias, or data_stream
510 Because you can pass search patterns and aliases as name, each element comes back
511 as an array:
513 {'indices': [], 'aliases': [], 'data_streams': []}
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'])
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')
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