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
« 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
6 # Importing bliss patches gevent which in turn patches python
7except ImportError:
8 from gevent.monkey import patch_all
10 patch_all(thread=False)
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
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
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
48import logging
50logger = logging.getLogger(__name__)
51DAIQUIRI_ROOT = os.path.dirname(daiquiri.__file__)
54def init_server(
55 options: ServerOptions,
56 testing: bool = False,
57):
58 """Instantiate a Flask application
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)
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
91 config = utils.ConfigDict("app.yml")
93 if not config.get("versions"):
94 config["versions"] = []
96 # Allow CLI to override implementors
97 if options.implementors:
98 config["implementors"] = options.implementors
100 log.start(config=config)
101 logger.info(f"Starting daiquiri version: {daiquiri.__version__}")
103 app = Flask(__name__, static_folder=static_folder, static_url_path="/")
105 app.config["APISPEC_FORMAT_RESPONSE"] = None
106 app.config["APISPEC_TITLE"] = "daiquiri"
108 security_definitions = {
109 "bearer": {"type": "apiKey", "in": "header", "name": "Authorization"}
110 }
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 )
120 if not config.get("swagger"):
121 app.config["APISPEC_SWAGGER_URL"] = None
122 app.config["APISPEC_SWAGGER_UI_URL"] = None
124 docs = FlaskApiSpec(app)
126 sio_kws = {}
127 if config["cors"]:
128 from flask_cors import CORS
130 sio_kws["cors_allowed_origins"] = "*"
131 CORS(app)
133 app.config["SECRET_KEY"] = config["iosecret"]
134 socketio = SocketIO(app, **sio_kws)
136 log.init_sio(socketio)
138 @app.errorhandler(404)
139 def page_not_found(e):
140 return jsonify(error=str(e)), 404
142 @app.errorhandler(405)
143 def method_not_allowed(e):
144 return jsonify(message="The method is not allowed for the requested URL."), 405
146 # @app.errorhandler(422)
147 # def unprocessable(e):
148 # return jsonify(error=str(e)), 422
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)
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)
163 Authenticator(config=config, app=app, session=ses, docs=docs, schema=schema)
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
172 BlissSession(config["controls_session_name"])
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
180 queue = Queue(
181 config=config,
182 app=app,
183 session=ses,
184 docs=docs,
185 socketio=socketio,
186 schema=schema,
187 )
189 Layout(
190 config, app=app, session=ses, docs=docs, socketio=socketio, schema=schema
191 )
193 metadata = MetaData(
194 config, app=app, session=ses, docs=docs, socketio=socketio, schema=schema
195 ).init()
196 ses.set_metadata(metadata)
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()
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 )
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)
229 exit()
231 if config["debug"] is True:
232 app.debug = True
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
242 def close():
243 components.close()
245 app.close = close
247 @app.route("/manifest.json")
248 def manifest():
249 return app.send_static_file("manifest.json")
251 @app.route("/meta.json")
252 @nocache
253 def meta():
254 return app.send_static_file("meta.json")
256 @app.route("/favicon.ico")
257 def favicon():
258 return app.send_static_file("favicon.ico")
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)
264 @app.route("/resources/<path:path>")
265 def resources(path):
266 """Serve a resource file from static_resources_folder if it exists.
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)
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")
281 took = round(time.time() - start, 2)
282 logger.info(f"Server ready, startup took {took}s", extra={"startup_time": took})
284 return app, socketio
287def get_certs_file_path(resource):
288 """Get a file path from a resource path
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)
300def get_ssl_context(options: ServerOptions):
301 """Build an ssl context for serving over HTTPS.
303 The options get the priority on the configuration.
305 The configurations should be dropped at some point for security.
306 """
307 config = utils.ConfigDict("app.yml")
309 use_ssl = options.ssl or config.get("ssl", False)
310 if not use_ssl:
311 return None
313 import ssl
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"])
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
325 return context
328def run_server(options: ServerOptions):
329 """Runs REST server
331 :param int port:
332 :param kwargs: see `init_server`
333 """
334 app, socketio = init_server(options)
336 server_args = {}
337 context = get_ssl_context(options)
338 if context:
339 server_args["ssl_context"] = context
341 socketio.run(app, host="0.0.0.0", port=options.port, **server_args) # nosec
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('_', '-')}"]
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 )
367def parse_options() -> ServerOptions:
368 """Parse the server options.
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()
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)
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})
395 options = ServerOptions(**merged_args)
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)
401 return options
404def main():
405 """Runs REST server with CLI configuration"""
406 options = parse_options()
407 run_server(options)
410if __name__ == "__main__":
411 main()