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

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 .fix_static_path_impl import recurse_directory 

29from .utils import none_to_default 

30 

31LOGGER: Final = logging.getLogger(__name__) 

32 

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

34 "component_templates", "index_templates", "ingest_pipelines" 

35] 

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

37 "ingest_pipelines", 

38 "component_templates", 

39 "index_templates", 

40) 

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 get: Callable[..., Awaitable[ObjectApiResponse[object]]] 

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

90 

91 if what == "component_templates": 

92 get = es.cluster.get_component_template 

93 put = es.cluster.put_component_template 

94 elif what == "index_templates": 

95 get = es.indices.get_index_template 

96 put = es.indices.put_index_template 

97 elif what == "ingest_pipelines": 

98 get = es.ingest.get_pipeline 

99 put = es.ingest.put_pipeline 

100 else: 

101 raise AssertionError() 

102 

103 try: 

104 if what == "ingest_pipelines": 

105 current = await get(id=name) 

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

107 else: 

108 current = await get( 

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

110 ) 

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

112 except NotFoundError: 

113 current_version = 0 

114 

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

116 if what == "ingest_pipelines": 

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

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

119 

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

121 LOGGER.warning( 

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

123 path, 

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

125 current_version, 

126 ) 

127 

128 return None 

129 

130 

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

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

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

134 from elastic_transport.client_utils import DEFAULT, DefaultType 

135 

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

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

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

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

140 ) 

141 

142 class Kwargs(TypedDict): 

143 """Kwargs of AsyncElasticsearch constructor.""" 

144 

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

146 cloud_id: None | str 

147 verify_certs: bool 

148 api_key: None | str 

149 bearer_auth: None | str 

150 client_cert: str | DefaultType 

151 client_key: str | DefaultType 

152 retry_on_timeout: bool | DefaultType 

153 

154 kwargs: Kwargs = { 

155 "hosts": ( 

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

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

158 else None 

159 ), 

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

161 "verify_certs": config.getboolean( 

162 "ELASTICSEARCH", "VERIFY_CERTS", fallback=True 

163 ), 

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

165 "bearer_auth": config.get( 

166 "ELASTICSEARCH", "BEARER_AUTH", fallback=None 

167 ), 

168 "client_cert": none_to_default( 

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

170 ), 

171 "client_key": none_to_default( 

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

173 ), 

174 "retry_on_timeout": none_to_default( 

175 config.getboolean( 

176 "ELASTICSEARCH", "RETRY_ON_TIMEOUT", fallback=None 

177 ), 

178 DEFAULT, 

179 ), 

180 } 

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

182 app.settings["ELASTICSEARCH"] = None 

183 return None 

184 elasticsearch = AsyncElasticsearch( 

185 basic_auth=( 

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

187 ), 

188 ca_certs=CA_BUNDLE_PATH, 

189 **kwargs, 

190 ) 

191 app.settings["ELASTICSEARCH"] = elasticsearch 

192 return elasticsearch