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

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.""" 

14from __future__ import annotations 

15 

16import asyncio 

17import logging 

18from collections.abc import Awaitable, Callable 

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

20 

21import orjson 

22from elastic_transport import ObjectApiResponse 

23from elasticsearch import AsyncElasticsearch, NotFoundError 

24from tornado.web import Application 

25 

26from .. import CA_BUNDLE_PATH, DIR 

27from .better_config_parser import BetterConfigParser 

28from .utils import none_to_default, recurse_directory 

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) 

40 

41 

42async def setup_elasticsearch_configs( 

43 elasticsearch: AsyncElasticsearch, 

44 prefix: str, 

45) -> None: 

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

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

48 

49 for i in range(3): 

50 spam = [] 

51 

52 what: ES_WHAT_LITERAL = ES_WHAT_LITERALS[i] 

53 

54 base_path = DIR / "elasticsearch" / what 

55 

56 for rel_path in recurse_directory( 

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

58 ): 

59 path = base_path / rel_path 

60 if not path.is_file(): 

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

62 continue 

63 

64 body = orjson.loads( 

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

66 ) 

67 

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

69 

70 spam.append( 

71 setup_elasticsearch_config( 

72 elasticsearch, what, body, name, rel_path 

73 ) 

74 ) 

75 

76 await asyncio.gather(*spam) 

77 

78 

79async def setup_elasticsearch_config( 

80 es: AsyncElasticsearch, 

81 what: ES_WHAT_LITERAL, 

82 body: dict[str, object], 

83 name: str, 

84 path: str = "<unknown>", 

85) -> None | ObjectApiResponse[object]: 

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

87 get: Callable[..., Awaitable[ObjectApiResponse[object]]] 

88 put: Callable[..., Awaitable[ObjectApiResponse[object]]] 

89 

90 if what == "component_templates": 

91 get = es.cluster.get_component_template 

92 put = es.cluster.put_component_template 

93 elif what == "index_templates": 

94 get = es.indices.get_index_template 

95 put = es.indices.put_index_template 

96 elif what == "ingest_pipelines": 

97 get = es.ingest.get_pipeline 

98 put = es.ingest.put_pipeline 

99 else: 

100 raise AssertionError() 

101 

102 try: 

103 if what == "ingest_pipelines": 

104 current = await get(id=name) 

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

106 else: 

107 current = await get( 

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

109 ) 

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

111 except NotFoundError: 

112 current_version = 0 

113 

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

115 if what == "ingest_pipelines": 

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

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

118 

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

120 LOGGER.warning( 

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

122 path, 

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

124 current_version, 

125 ) 

126 

127 return None 

128 

129 

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

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

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

133 from elastic_transport.client_utils import DEFAULT, DefaultType 

134 

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

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

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

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

139 ) 

140 

141 class Kwargs(TypedDict): 

142 """Kwargs of AsyncElasticsearch constructor.""" 

143 

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

145 cloud_id: None | str 

146 verify_certs: bool 

147 api_key: None | str 

148 bearer_auth: None | str 

149 client_cert: str | DefaultType 

150 client_key: str | DefaultType 

151 retry_on_timeout: bool | DefaultType 

152 

153 kwargs: Kwargs = { 

154 "hosts": ( 

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

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

157 else None 

158 ), 

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

160 "verify_certs": config.getboolean( 

161 "ELASTICSEARCH", "VERIFY_CERTS", fallback=True 

162 ), 

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

164 "bearer_auth": config.get( 

165 "ELASTICSEARCH", "BEARER_AUTH", fallback=None 

166 ), 

167 "client_cert": none_to_default( 

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

169 ), 

170 "client_key": none_to_default( 

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

172 ), 

173 "retry_on_timeout": none_to_default( 

174 config.getboolean( 

175 "ELASTICSEARCH", "RETRY_ON_TIMEOUT", fallback=None 

176 ), 

177 DEFAULT, 

178 ), 

179 } 

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

181 app.settings["ELASTICSEARCH"] = None 

182 return None 

183 elasticsearch = AsyncElasticsearch( 

184 basic_auth=( 

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

186 ), 

187 ca_certs=CA_BUNDLE_PATH, 

188 **kwargs, 

189 ) 

190 app.settings["ELASTICSEARCH"] = elasticsearch 

191 return elasticsearch