Coverage for /opt/conda/envs/apienv/lib/python3.10/site-packages/daiquiri/cli/server.py: 50%

216 statements  

« prev     ^ index     » next       coverage.py v7.6.10, created at 2025-02-06 02:13 +0000

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3try: 

4 import bliss # noqa F401 

5 

6 # Importing bliss patches gevent which in turn patches python 

7except ImportError: 

8 from gevent.monkey import patch_all 

9 

10 patch_all(thread=False) 

11 

12# TODO: Horrible hack to avoid: 

13# ImportError: dlopen: cannot load any more object with static TLS 

14# https://github.com/pytorch/pytorch/issues/2575 

15# https://github.com/scikit-learn/scikit-learn/issues/14485 

16from silx.math import colormap # noqa F401 

17 

18import time 

19import os 

20import json 

21import argparse 

22import pydantic 

23from ruamel.yaml import YAML 

24from flask import Flask, jsonify 

25from flask_socketio import SocketIO 

26from flask_apispec import FlaskApiSpec 

27from flask_restful import abort 

28from flask import send_from_directory 

29from webargs.flaskparser import parser 

30from apispec import APISpec 

31from apispec.ext.marshmallow import MarshmallowPlugin 

32 

33import daiquiri 

34from daiquiri.core.authenticator import Authenticator 

35from daiquiri.core.session import Session 

36from daiquiri.core.queue import Queue 

37from daiquiri.core.metadata import MetaData 

38from daiquiri.core.components import Components 

39from daiquiri.core.hardware import Hardware 

40from daiquiri.core.schema import Schema 

41from daiquiri.core.layout import Layout 

42from daiquiri.core.saving import Saving 

43from daiquiri.resources import utils 

44from daiquiri.core.logging import log 

45from daiquiri.core.responses import nocache 

46from daiquiri.core.options import ServerOptions 

47 

48import logging 

49 

50logger = logging.getLogger(__name__) 

51DAIQUIRI_ROOT = os.path.dirname(daiquiri.__file__) 

52 

53 

54def init_server( 

55 options: ServerOptions, 

56 testing: bool = False, 

57): 

58 """Instantiate a Flask application 

59 

60 :param options: Options use to setup the server 

61 :param testing: Attaches core components to the `app` so that can be retrieved in tests 

62 :returns: Flask, SocketIO 

63 """ 

64 start = time.time() 

65 for resource_folder in reversed(options.resource_folders): 

66 if not resource_folder: 

67 continue 

68 if not os.path.isdir(resource_folder): 

69 raise ValueError(f"Resource folder '{resource_folder}' does not exist") 

70 utils.add_resource_root(resource_folder) 

71 logger.info("Added resource folder: %s", resource_folder) 

72 

73 static_folder = options.static_folder 

74 if static_folder.startswith("static."): 

75 provider = utils.get_resource_provider() 

76 try: 

77 static_folder = provider.get_resource_path(static_folder, "") 

78 except utils.ResourceNotAvailable: 

79 raise ValueError( 

80 f"Static resource '{static_folder}' does not exist" 

81 ) from None 

82 if not os.path.isdir(static_folder): 

83 raise ValueError(f"Static folder '{static_folder}' does not exist") 

84 if options.hardware_folder: 

85 if not os.path.isdir(options.hardware_folder): 

86 raise ValueError( 

87 f"Hardware folder '{options.hardware_folder}' does not exist" 

88 ) 

89 os.environ["HWR_ROOT"] = options.hardware_folder 

90 

91 config = utils.ConfigDict("app.yml") 

92 

93 if not config.get("versions"): 

94 config["versions"] = [] 

95 

96 # Allow CLI to override implementors 

97 if options.implementors: 

98 config["implementors"] = options.implementors 

99 

100 log.start(config=config) 

101 logger.info(f"Starting daiquiri version: {daiquiri.__version__}") 

102 

103 app = Flask(__name__, static_folder=static_folder, static_url_path="/") 

104 

105 app.config["APISPEC_FORMAT_RESPONSE"] = None 

106 app.config["APISPEC_TITLE"] = "daiquiri" 

107 

108 security_definitions = { 

109 "bearer": {"type": "apiKey", "in": "header", "name": "Authorization"} 

110 } 

111 

112 app.config["APISPEC_SPEC"] = APISpec( 

113 title="daiquiri", 

114 version="v1", 

115 openapi_version="2.0", 

116 plugins=[MarshmallowPlugin()], 

117 securityDefinitions=security_definitions, 

118 ) 

119 

120 if not config.get("swagger"): 

121 app.config["APISPEC_SWAGGER_URL"] = None 

122 app.config["APISPEC_SWAGGER_UI_URL"] = None 

123 

124 docs = FlaskApiSpec(app) 

125 

126 sio_kws = {} 

127 if config["cors"]: 

128 from flask_cors import CORS 

129 

130 sio_kws["cors_allowed_origins"] = "*" 

131 CORS(app) 

132 

133 app.config["SECRET_KEY"] = config["iosecret"] 

134 socketio = SocketIO(app, **sio_kws) 

135 

136 log.init_sio(socketio) 

137 

138 @app.errorhandler(404) 

139 def page_not_found(e): 

140 return jsonify(error=str(e)), 404 

141 

142 @app.errorhandler(405) 

143 def method_not_allowed(e): 

144 return jsonify(message="The method is not allowed for the requested URL."), 405 

145 

146 # @app.errorhandler(422) 

147 # def unprocessable(e): 

148 # return jsonify(error=str(e)), 422 

149 

150 @parser.error_handler 

151 def handle_request_parsing_error( 

152 err, req, schema, error_status_code, error_headers 

153 ): 

154 abort(422, description=err.messages) 

155 

156 if not app.debug or os.environ.get("WERKZEUG_RUN_MAIN") == "true": 

157 schema = Schema(app=app, docs=docs, socketio=socketio) 

158 ses = Session( 

159 config=config, app=app, docs=docs, socketio=socketio, schema=schema 

160 ) 

161 schema.set_session(ses) 

162 

163 Authenticator(config=config, app=app, session=ses, docs=docs, schema=schema) 

164 

165 # TODO 

166 # This is BLISS specific, move to where bliss is handled ? 

167 if config.get("controls_session_type", None) == "bliss": 

168 # This is not very elegant but a solution until this section 

169 # have been moved to an appropriate place 

170 from daiquiri.core.hardware.bliss.session import BlissSession 

171 

172 BlissSession(config["controls_session_name"]) 

173 

174 hardware = Hardware( 

175 base_config=config, app=app, socketio=socketio, docs=docs, schema=schema 

176 ) 

177 # TODO: implement get_core_component API 

178 app.hardware = hardware 

179 

180 queue = Queue( 

181 config=config, 

182 app=app, 

183 session=ses, 

184 docs=docs, 

185 socketio=socketio, 

186 schema=schema, 

187 ) 

188 

189 Layout( 

190 config, app=app, session=ses, docs=docs, socketio=socketio, schema=schema 

191 ) 

192 

193 metadata = MetaData( 

194 config, app=app, session=ses, docs=docs, socketio=socketio, schema=schema 

195 ).init() 

196 ses.set_metadata(metadata) 

197 

198 saving = Saving( 

199 config, 

200 app=app, 

201 session=ses, 

202 docs=docs, 

203 socketio=socketio, 

204 schema=schema, 

205 metadata=metadata, 

206 ).init() 

207 

208 components = Components( 

209 base_config=config, 

210 app=app, 

211 socketio=socketio, 

212 docs=docs, 

213 schema=schema, 

214 hardware=hardware, 

215 session=ses, 

216 metadata=metadata, 

217 queue=queue, 

218 saving=saving, 

219 ) 

220 

221 if options.save_spec_file is not None: 

222 logger.info("Writing API spec and exiting") 

223 save_spec_dir = os.path.dirname(options.save_spec_file) 

224 if not os.path.exists(save_spec_dir): 

225 os.mkdir(save_spec_dir) 

226 with open(options.save_spec_file, "w") as spec: 

227 json.dump(docs.spec.to_dict(), spec) 

228 

229 exit() 

230 

231 if config["debug"] is True: 

232 app.debug = True 

233 

234 if testing: 

235 app.queue = queue 

236 app.metadata = metadata 

237 app.components = components 

238 app.session = ses 

239 app.socketio = socketio 

240 app.saving = saving 

241 

242 def close(): 

243 components.close() 

244 

245 app.close = close 

246 

247 @app.route("/manifest.json") 

248 def manifest(): 

249 return app.send_static_file("manifest.json") 

250 

251 @app.route("/meta.json") 

252 @nocache 

253 def meta(): 

254 return app.send_static_file("meta.json") 

255 

256 @app.route("/favicon.ico") 

257 def favicon(): 

258 return app.send_static_file("favicon.ico") 

259 

260 if options.static_resources_folder: 

261 default_static_resources = os.path.join(static_folder, "resources") 

262 static_resources_folder = os.path.abspath(options.static_resources_folder) 

263 

264 @app.route("/resources/<path:path>") 

265 def resources(path): 

266 """Serve a resource file from static_resources_folder if it exists. 

267 

268 Else fall back to the original static directory. 

269 """ 

270 if os.path.isfile(os.path.join(static_resources_folder, path)): 

271 return send_from_directory(static_resources_folder, path) 

272 return send_from_directory(default_static_resources, path) 

273 

274 @app.route("/", defaults={"path": ""}) 

275 @app.route("/<string:path>") 

276 @app.route("/<path:path>") 

277 @nocache 

278 def index(path): 

279 return app.send_static_file("index.html") 

280 

281 took = round(time.time() - start, 2) 

282 logger.info(f"Server ready, startup took {took}s", extra={"startup_time": took}) 

283 

284 return app, socketio 

285 

286 

287def get_certs_file_path(resource): 

288 """Get a file path from a resource path 

289 

290 This can be: 

291 - foobar.txt # read file from resource provides "certs" 

292 - /etc/foobar.txt # read file from absolute path 

293 """ 

294 if resource.startswith("/"): 

295 return os.path.abspath(resource) 

296 provider = utils.get_resource_provider() 

297 return provider.get_resource_path("certs", resource) 

298 

299 

300def get_ssl_context(options: ServerOptions): 

301 """Build an ssl context for serving over HTTPS. 

302 

303 The options get the priority on the configuration. 

304 

305 The configurations should be dropped at some point for security. 

306 """ 

307 config = utils.ConfigDict("app.yml") 

308 

309 use_ssl = options.ssl or config.get("ssl", False) 

310 if not use_ssl: 

311 return None 

312 

313 import ssl 

314 

315 crt = get_certs_file_path(options.ssl_cert or config["ssl_cert"]) 

316 key = get_certs_file_path(options.ssl_key or config["ssl_key"]) 

317 

318 context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2) 

319 try: 

320 context.load_cert_chain(crt, key) 

321 except Exception: 

322 logger.exception("Could not load certificate chain", exc_info=True) 

323 raise 

324 

325 return context 

326 

327 

328def run_server(options: ServerOptions): 

329 """Runs REST server 

330 

331 :param int port: 

332 :param kwargs: see `init_server` 

333 """ 

334 app, socketio = init_server(options) 

335 

336 server_args = {} 

337 context = get_ssl_context(options) 

338 if context: 

339 server_args["ssl_context"] = context 

340 

341 socketio.run(app, host="0.0.0.0", port=options.port, **server_args) # nosec 

342 

343 

344def add_pydantic_model(parser: argparse.ArgumentParser, model: pydantic.BaseModel): 

345 "Add Pydantic model to an ArgumentParser" 

346 fields = model.__fields__ 

347 for name, field in fields.items(): 

348 extra = field.json_schema_extra or {} 

349 name_or_flags = [] 

350 if "flag" in extra: 

351 name_or_flags += [extra["argparse_flag"]] 

352 if "argname" in extra: 

353 name_or_flags += [extra["argparse_name"]] 

354 else: 

355 name_or_flags += [f"--{name.replace('_', '-')}"] 

356 

357 parser.add_argument( 

358 *name_or_flags, 

359 dest=name, 

360 nargs=extra.get("argparse_nargs"), 

361 type=extra.get("argparse_type") or field.annotation, 

362 default=None, # let pydantic handle the default later 

363 help=field.description, 

364 ) 

365 

366 

367def parse_options() -> ServerOptions: 

368 """Parse the server options. 

369 

370 Read the command line arguments and the optional configuration file. 

371 """ 

372 parser = argparse.ArgumentParser(description="REST server for a beamline GUI") 

373 parser.add_argument( 

374 "-c", 

375 "--config", 

376 type=str, 

377 default=None, 

378 help="File/resource which is used to setup the server (instead of the command line argument)", 

379 ) 

380 add_pydantic_model(parser, ServerOptions) 

381 cmd_args = parser.parse_args() 

382 

383 config_args = {} 

384 if cmd_args.config: 

385 logging.info("Read config file: %s", cmd_args.config) 

386 yaml = YAML(typ="safe") 

387 with open(cmd_args.config, mode="rt") as f: 

388 config_args = yaml.load(f) 

389 

390 merged_args = {} 

391 merged_args.update(config_args) 

392 # Only override command line arguments which are set 

393 merged_args.update({k: v for k, v in vars(cmd_args).items() if v is not None}) 

394 

395 options = ServerOptions(**merged_args) 

396 

397 # Sounds like flask except an absolute path 

398 if options.static_folder is not None: 

399 options.static_folder = os.path.abspath(options.static_folder) 

400 

401 return options 

402 

403 

404def main(): 

405 """Runs REST server with CLI configuration""" 

406 options = parse_options() 

407 run_server(options) 

408 

409 

410if __name__ == "__main__": 

411 main()