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

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 an_website.utils.elastic_transport_async_http_node import TornadoAsyncNode 

26 

27from .. import CA_BUNDLE_PATH, DIR 

28from .better_config_parser import BetterConfigParser 

29from .fix_static_path_impl import recurse_directory 

30from .utils import none_to_default 

31 

32LOGGER: Final = logging.getLogger(__name__) 

33 

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

35 "component_templates", "index_templates", "ingest_pipelines" 

36] 

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

38 "ingest_pipelines", 

39 "component_templates", 

40 "index_templates", 

41) 

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

43 

44 

45async def setup_elasticsearch_configs( 

46 elasticsearch: AsyncElasticsearch, 

47 prefix: str, 

48) -> None: 

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

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

51 

52 for i in range(3): 

53 spam = [] 

54 

55 what: ES_WHAT_LITERAL = ES_WHAT_LITERALS[i] 

56 

57 base_path = DIR / "elasticsearch" / what 

58 

59 for rel_path in recurse_directory( 

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

61 ): 

62 path = base_path / rel_path 

63 if not path.is_file(): 

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

65 continue 

66 

67 body = orjson.loads( 

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

69 ) 

70 

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

72 

73 spam.append( 

74 setup_elasticsearch_config( 

75 elasticsearch, what, body, name, rel_path 

76 ) 

77 ) 

78 

79 await asyncio.gather(*spam) 

80 

81 

82async def setup_elasticsearch_config( 

83 es: AsyncElasticsearch, 

84 what: ES_WHAT_LITERAL, 

85 body: dict[str, object], 

86 name: str, 

87 path: str = "<unknown>", 

88) -> None | ObjectApiResponse[object]: 

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

90 if what == "component_templates": 

91 get: AnyArgsAsyncMethod = es.cluster.get_component_template 

92 put: AnyArgsAsyncMethod = 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 http_compress=True, 

189 node_class=TornadoAsyncNode, 

190 **kwargs, 

191 ) 

192 app.settings["ELASTICSEARCH"] = elasticsearch 

193 return elasticsearch