Coverage for an_website / utils / elasticsearch_setup.py: 40.299%

1# This program is free software: you can redistribute it and/or modify 

2# it under the terms of the GNU Affero General Public License as 

3# published by the Free Software Foundation, either version 3 of the 

4# License, or (at your option) any later version. 

5# 

6# This program is distributed in the hope that it will be useful, 

7# but WITHOUT ANY WARRANTY; without even the implied warranty of 

8# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

9# GNU Affero General Public License for more details. 

10# 

11# You should have received a copy of the GNU Affero General Public License 

12# along with this program. If not, see <https://www.gnu.org/licenses/>. 

13"""Functions for setting up Elasticsearch.""" 

14 

15import asyncio 

16import logging 

17from collections.abc import Awaitable, Callable 

18from typing import Final, Literal, TypeAlias, TypedDict, cast 

19 

20import orjson 

21from elastic_transport import ObjectApiResponse 

22from elasticsearch import AsyncElasticsearch, NotFoundError 

23from tornado.web import Application 

24 

25from .. import CA_BUNDLE_PATH, DIR 

26from .better_config_parser import BetterConfigParser 

27from .fix_static_path_impl import recurse_directory 

28from .utils import none_to_default 

29 

30LOGGER: Final = logging.getLogger(__name__) 

31 

32ES_WHAT_LITERAL: TypeAlias = Literal[ # pylint: disable=invalid-name 

33 "component_templates", "index_templates", "ingest_pipelines" 

34] 

35ES_WHAT_LITERALS: tuple[ES_WHAT_LITERAL, ...] = ( 

36 "ingest_pipelines", 

37 "component_templates", 

38 "index_templates", 

39) 

40type AnyArgsAsyncMethod = Callable[..., Awaitable[ObjectApiResponse[object]]] 

41 

42 

43async def setup_elasticsearch_configs( 

44 elasticsearch: AsyncElasticsearch, 

45 prefix: str, 

46) -> None: 

47 """Setup Elasticsearch configs.""" # noqa: D401 

48 spam: list[Awaitable[None | ObjectApiResponse[object]]] 

49 

50 for i in range(3): 

51 spam = [] 

52 

53 what: ES_WHAT_LITERAL = ES_WHAT_LITERALS[i] 

54 

55 base_path = DIR / "elasticsearch" / what 

56 

57 for rel_path in recurse_directory( 

58 base_path, lambda path: path.name.endswith(".json") 

59 ): 

60 path = base_path / rel_path 

61 if not path.is_file(): 

62 LOGGER.warning("%s is not a file", path) 

63 continue 

64 

65 body = orjson.loads( 

66 path.read_bytes().replace(b"{prefix}", prefix.encode("ASCII")) 

67 ) 

68 

69 name = f"{prefix}-{rel_path[:-5].replace('/', '-')}" 

70 

71 spam.append( 

72 setup_elasticsearch_config( 

73 elasticsearch, what, body, name, rel_path 

74 ) 

75 ) 

76 

77 await asyncio.gather(*spam) 

78 

79 

80async def setup_elasticsearch_config( 

81 es: AsyncElasticsearch, 

82 what: ES_WHAT_LITERAL, 

83 body: dict[str, object], 

84 name: str, 

85 path: str = "<unknown>", 

86) -> None | ObjectApiResponse[object]: 

87 """Setup Elasticsearch config.""" # noqa: D401 

88 if what == "component_templates": 

89 get: AnyArgsAsyncMethod = es.cluster.get_component_template 

90 put: AnyArgsAsyncMethod = es.cluster.put_component_template 

91 elif what == "index_templates": 

92 get = es.indices.get_index_template 

93 put = es.indices.put_index_template 

94 elif what == "ingest_pipelines": 

95 get = es.ingest.get_pipeline 

96 put = es.ingest.put_pipeline 

97 else: 

98 raise AssertionError() 

99 

100 try: 

101 if what == "ingest_pipelines": 

102 current = await get(id=name) 

103 current_version = current[name].get("version", 1) 

104 else: 

105 current = await get( 

106 name=name, filter_path=f"{what}.name,{what}.version" 

107 ) 

108 current_version = current[what][0].get("version", 1) 

109 except NotFoundError: 

110 current_version = 0 

111 

112 if current_version < body.get("version", 1): 

113 if what == "ingest_pipelines": 

114 return await put(id=name, body=body) 

115 return await put(name=name, body=body) 

116 

117 if current_version > body.get("version", 1): 

118 LOGGER.warning( 

119 "%s has version %s. The version in Elasticsearch is %s!", 

120 path, 

121 body.get("version", 1), 

122 current_version, 

123 ) 

124 

125 return None 

126 

127 

128def setup_elasticsearch(app: Application) -> None | AsyncElasticsearch: 

129 """Setup Elasticsearch.""" # noqa: D401 

130 # pylint: disable-next=import-outside-toplevel 

131 from elastic_transport.client_utils import DEFAULT, DefaultType 

132 

133 config: BetterConfigParser = app.settings["CONFIG"] 

134 basic_auth: tuple[str | None, str | None] = ( 

135 config.get("ELASTICSEARCH", "USERNAME", fallback=None), 

136 config.get("ELASTICSEARCH", "PASSWORD", fallback=None), 

137 ) 

138 

139 class Kwargs(TypedDict): 

140 """Kwargs of AsyncElasticsearch constructor.""" 

141 

142 hosts: tuple[str, ...] | None 

143 cloud_id: None | str 

144 verify_certs: bool 

145 api_key: None | str 

146 bearer_auth: None | str 

147 client_cert: str | DefaultType 

148 client_key: str | DefaultType 

149 retry_on_timeout: bool | DefaultType 

150 

151 kwargs: Kwargs = { 

152 "hosts": ( 

153 tuple(config.getset("ELASTICSEARCH", "HOSTS")) 

154 if config.has_option("ELASTICSEARCH", "HOSTS") 

155 else None 

156 ), 

157 "cloud_id": config.get("ELASTICSEARCH", "CLOUD_ID", fallback=None), 

158 "verify_certs": config.getboolean( 

159 "ELASTICSEARCH", "VERIFY_CERTS", fallback=True 

160 ), 

161 "api_key": config.get("ELASTICSEARCH", "API_KEY", fallback=None), 

162 "bearer_auth": config.get( 

163 "ELASTICSEARCH", "BEARER_AUTH", fallback=None 

164 ), 

165 "client_cert": none_to_default( 

166 config.get("ELASTICSEARCH", "CLIENT_CERT", fallback=None), DEFAULT 

167 ), 

168 "client_key": none_to_default( 

169 config.get("ELASTICSEARCH", "CLIENT_KEY", fallback=None), DEFAULT 

170 ), 

171 "retry_on_timeout": none_to_default( 

172 config.getboolean( 

173 "ELASTICSEARCH", "RETRY_ON_TIMEOUT", fallback=None 

174 ), 

175 DEFAULT, 

176 ), 

177 } 

178 if not config.getboolean("ELASTICSEARCH", "ENABLED", fallback=False): 

179 app.settings["ELASTICSEARCH"] = None 

180 return None 

181 elasticsearch = AsyncElasticsearch( 

182 basic_auth=( 

183 None if None in basic_auth else cast(tuple[str, str], basic_auth) 

184 ), 

185 ca_certs=CA_BUNDLE_PATH, 

186 **kwargs, 

187 ) 

188 app.settings["ELASTICSEARCH"] = elasticsearch 

189 return elasticsearch