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

86 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-05-02 23:00 -0600

1"""TestPlan Class Definition""" 

2 

3import typing as t 

4import logging 

5from dotmap import DotMap 

6from es_testbed.defaults import TESTPLAN 

7from es_testbed.helpers.utils import build_ilm_policy, prettystr, randomstr 

8 

9logger = logging.getLogger('es_testbed.PlanBuilder') 

10 

11 

12class PlanBuilder: 

13 """Plan builder class""" 

14 

15 def __init__( 

16 self, 

17 settings: t.Dict = None, 

18 default_entities: bool = True, 

19 autobuild: t.Optional[bool] = True, 

20 ): 

21 self.default_entities = default_entities 

22 if settings is None: 

23 raise ValueError('Must provide a settings dictionary') 

24 self.settings = settings 

25 self._plan = DotMap(TESTPLAN) 

26 logger.debug('INITIAL PLAN: %s', prettystr(self._plan)) 

27 self._plan.cleanup = 'UNKNOWN' # Future use? 

28 if autobuild: 

29 self.setup() 

30 

31 # ## Example settings 

32 # settings={ 

33 # 'type': 'indices', # Default is indices? Or should it be data_streams? 

34 # 'prefix': 'es-testbed', # Provide this value as a default 

35 # 'rollover_alias': False, # Only respected if 'type' == 'indices'. 

36 # # Will rollover after creation and filling 1st 

37 # # If True, will be overridden to value of alias 

38 # # If False, will be overridden with None 

39 # 'uniq': 'my-unique-str', # If not provided, randomstr() 

40 # 'repository': # Only used for cold/frozen tier for snapshots 

41 # 'ilm': { # All of these ILM values are defaults 

42 # 'enabled': False, 

43 # 'tiers': ['hot', 'delete'], 

44 # 'forcemerge': False, 

45 # 'max_num_segments': 1, 

46 # } 

47 # 

48 # # If these keys aren't specified per entity, then all entities will get this 

49 # # treatment 

50 # # EXCEPT for the is_write_index for aliases and data_streams 

51 # 

52 # 'defaults': { 

53 # 'entity_count': 3, 

54 # 'docs': 10, 

55 # 'match': True, 

56 # 'searchable': tier... 

57 # } 

58 # 

59 # # Manually specifying entities makes sense for individual indices, but not so 

60 # # much for 

61 # # alias-backed indices or data_streams 

62 # 'entities': [ 

63 # { 

64 # 'docs': 10, 

65 # 'match': True, 

66 # 'searchable': 'frozen' 

67 # }, 

68 # { 

69 # 'docs': 10, 

70 # 'match': False, 

71 # 'searchable': 'cold', 

72 # }, 

73 # { 

74 # 'docs': 10, 

75 # 'match': True, 

76 # 'searchable': 'hot' 

77 # }, 

78 # ] 

79 # } 

80 

81 @property 

82 def plan(self) -> DotMap: 

83 """Return the Plan""" 

84 return self._plan 

85 

86 def _create_lists(self) -> None: 

87 names = [ 

88 'indices', 

89 'data_stream', 

90 'snapshots', 

91 'ilm_policies', 

92 'index_templates', 

93 'component_templates', 

94 ] 

95 for name in names: 

96 self._plan[name] = [] 

97 

98 def add_entity( 

99 self, 

100 docs: t.Optional[int] = 10, 

101 match: t.Optional[bool] = True, 

102 searchable: t.Optional[str] = None, 

103 ) -> None: 

104 """Add an index or data_stream""" 

105 entity = DotMap({'docs': docs, 'match': match}) 

106 if searchable: 

107 entity.searchable = searchable 

108 self._plan.entities.append(entity) 

109 

110 def make_default_entities(self) -> None: 

111 """Loop through until all entities are created""" 

112 defs = TESTPLAN['defaults'] # Start with defaults 

113 if 'defaults' in self._plan: 

114 defs = self._plan.defaults 

115 kwargs = { 

116 'docs': defs['docs'], 

117 'match': defs['match'], 

118 'searchable': defs['searchable'], 

119 } 

120 self._plan.entities = [] 

121 for _ in range(0, defs['entity_count']): 

122 self.add_entity(**kwargs) 

123 logger.debug('Plan will create %s (backing) indices', len(self._plan.entities)) 

124 

125 def setup(self) -> None: 

126 """Do initial setup of the Plan DotMap""" 

127 self._plan.uniq = randomstr(length=8, lowercase=True) 

128 self._create_lists() 

129 self.update(self.settings) # Override with settings. 

130 self.update_rollover_alias() 

131 logger.debug('Rollover alias updated') 

132 self.update_ilm() 

133 if not self._plan.entities: 

134 if self.default_entities: 

135 self.make_default_entities() 

136 logger.debug('Test Plan: %s', prettystr(self._plan.toDict())) 

137 

138 def update(self, settings: t.Dict) -> None: 

139 """Update the Plan DotMap""" 

140 self._plan.update(**settings) 

141 

142 def update_ilm(self) -> None: 

143 """Update the ILM portion of the Plan DotMap""" 

144 setdefault = False 

145 if 'ilm' not in self._plan: 

146 logger.debug('key "ilm" is not in plan') 

147 setdefault = True 

148 if isinstance(self._plan.ilm, dict): 

149 _ = DotMap(self._plan.ilm) 

150 self._plan.ilm = _ 

151 if isinstance(self._plan.ilm, DotMap): 

152 if 'enabled' not in self._plan.ilm: 

153 # Override with defaults 

154 logger.debug( 

155 'plan.ilm does not have key "enabled". Overriding with defaults' 

156 ) 

157 setdefault = True 

158 elif isinstance(self._plan.ilm, bool): 

159 if self._plan.ilm: 

160 logger.warning( 

161 '"plan.ilm: True" is incorrect. Use plan.ilm.enabled: True' 

162 ) 

163 logger.debug('plan.ilm is boolean. Overriding with defaults') 

164 setdefault = True 

165 if setdefault: 

166 logger.debug('Setting defaults for ILM') 

167 self._plan.ilm = DotMap(TESTPLAN['ilm']) 

168 if self._plan.ilm.enabled: 

169 ilm = self._plan.ilm 

170 if not isinstance(self._plan.ilm.tiers, list): 

171 logger.error('Tiers is not a list!') 

172 self._plan.ilm.tiers = TESTPLAN['ilm']['tiers'] 

173 for entity in self._plan.entities: 

174 if 'searchable' in entity and entity['searchable'] is not None: 

175 if not entity['searchable'] in ilm.tiers: 

176 ilm.tiers.append(entity['searchable']) 

177 kwargs = { 

178 'tiers': ilm.tiers, 

179 'forcemerge': ilm.forcemerge, 

180 'max_num_segments': ilm.max_num_segments, 

181 'repository': self._plan.repository, 

182 } 

183 self._plan.ilm.policy = build_ilm_policy(**kwargs) 

184 

185 def update_rollover_alias(self) -> None: 

186 """Update the Rollover Alias value""" 

187 if self._plan.rollover_alias: 

188 self._plan.rollover_alias = f'{self._plan.prefix}-idx-{self._plan.uniq}' 

189 else: 

190 self._plan.rollover_alias = None