Replaced flask_restplus use of ujson with orjson. Fixed precision problem.

nomad/app/__init__.py

@@ -17,23 +17,55 @@ This module comprises the nomad@FAIRDI APIs. Currently there is NOMAD's official
 we will soon at the optimade api. The app module also servers documentation, gui, and
 alive.
 '''
-from flask import Flask, Blueprint, jsonify, url_for, abort, request
-from flask_restplus import Api
+from flask import Flask, Blueprint, jsonify, url_for, abort, request, make_response
+from flask_restplus import Api, representations
 from flask_cors import CORS
 from werkzeug.exceptions import HTTPException
 from werkzeug.wsgi import DispatcherMiddleware  # pylint: disable=E0611
 import os.path
 import random
 from structlog import BoundLogger
+import orjson
+import collections
+from mongoengine.base.datastructures import BaseList
 
 from nomad import config, utils as nomad_utils
 
-from .api import blueprint as api
-from .optimade import blueprint as optimade
-from .docs import blueprint as docs
+from .api import blueprint as api_blueprint, api
+from .optimade import blueprint as optimade_blueprint, api as optimade
+from .docs import blueprint as docs_blueprint
 from . import common
 
 
+__new_line = '\n'.encode()
+
+
+# replace the json implementation of flask_restplus
+def output_json(data, code, headers=None):
+    def default(data):
+        if isinstance(data, collections.OrderedDict):
+            return dict(data)
+
+        if isinstance(data, BaseList):
+            return list(data)
+
+        raise TypeError
+
+    # always end the json dumps with a new line
+    # see https://github.com/mitsuhiko/flask/pull/1262
+    dumped = orjson.dumps(
+        data, default=default,
+        option=orjson.OPT_INDENT_2 | orjson.OPT_NON_STR_KEYS) + __new_line
+
+    resp = make_response(dumped, code)
+    resp.headers.extend(headers or {})
+    return resp
+
+
+api.representation('application/json')(output_json)
+optimade.representation('application/json')(output_json)
+
+
 @property  # type: ignore
 def specs_url(self):
     '''
@@ -72,9 +104,9 @@ app.wsgi_app = DispatcherMiddleware(  # type: ignore
 
 CORS(app)
 
-app.register_blueprint(api, url_prefix='/api')
-app.register_blueprint(optimade, url_prefix='/optimade')
-app.register_blueprint(docs, url_prefix='/docs')
+app.register_blueprint(api_blueprint, url_prefix='/api')
+app.register_blueprint(optimade_blueprint, url_prefix='/optimade')
+app.register_blueprint(docs_blueprint, url_prefix='/docs')
 
 
 @app.errorhandler(Exception)

nomad/app/api/__init__.py

@@ -24,5 +24,5 @@ There is a separate documentation for the API endpoints from a client perspectiv
 .. automodule:: nomad.app.api.archive
 '''
 
-from .api import blueprint
+from .api import api, blueprint
 from . import info, auth, upload, repo, archive, raw, mirror, dataset

nomad/app/api/repo.py

@@ -622,8 +622,11 @@ _repo_quantities_search_request_parser.add_argument(
 _repo_quantities_search_request_parser.add_argument(
     'size', type=int, help='The max size of the returned values.')
 
-_repo_quantities_model = api.model('RepoQuantities', {
-    'quantities': fields.List(fields.Nested(_repo_quantity_model))
+_repo_quantities_model = api.model('RepoQuantitiesResponse', {
+    'quantities': fields.Nested(api.model('RepoQuantities', {
+        quantity: fields.List(fields.Nested(_repo_quantity_model))
+        for quantity in search_extension.search_quantities
+    }))
 })
 
 

nomad/app/optimade/__init__.py

@@ -19,6 +19,6 @@ The optimade implementation of NOMAD.
 from flask import Blueprint
 from flask_restplus import Api
 
-from .api import blueprint, url
+from .api import blueprint, url, api
 from .endpoints import CalculationList, Calculation
 from .filterparser import parse_filter

nomad/app/optimade/endpoints.py

@@ -144,7 +144,7 @@ class CalculationInfo(Resource):
                 for attr in OptimadeEntry.m_def.all_properties.values()},
             'formats': ['json'],
             'output_fields_by_format': {
-                'json': OptimadeEntry.m_def.all_properties.keys()}
+                'json': list(OptimadeEntry.m_def.all_properties.keys())}
         }
 
         return dict(

nomad/search.py

@@ -574,6 +574,7 @@ class SearchRequest:
         # statistics
         def get_metrics(bucket, code_runs):
             result = {}
+            # TODO optimize ... go through the buckets not the metrics
             for metric in metrics:
                 agg_name = 'metric:%s' % metric
                 if agg_name in bucket:
@@ -583,7 +584,7 @@ class SearchRequest:
 
         statistics_results = {
             quantity_name[11:]: {
-                bucket['key']: get_metrics(bucket, bucket['doc_count'])
+                str(bucket['key']): get_metrics(bucket, bucket['doc_count'])
                 for bucket in quantity['buckets']
             }
             for quantity_name, quantity in aggs.items()

requirements.txt

@@ -11,7 +11,7 @@ Pint
 matid
 
 # infrastructue related
-ujson<2.0.0
+orjson
 elasticsearch==6.4.0
 elasticsearch-dsl==6.4.0
 msgpack<0.6.0

tests/app/test_api.py

@@ -1084,7 +1084,9 @@ class TestRepo():
                 dict(quantities=['dft.system', 'atoms'], size=1), doseq=True),
             headers=test_user_auth)
         assert rv.status_code == 200
-        # TODO actual assertions
+        data = json.loads(rv.data)
+        assert 'dft.system' in data['quantities']
+        assert 'atoms' in data['quantities']
 
     @pytest.mark.parametrize('pid_or_handle, with_login, success', [
         ('2', True, True), ('2', False, True),