REST API: POST endpoint for QueryBuilder queryhelp JSON payloads

.ci/polish/cli.py

@@ -102,26 +102,27 @@ def launch(expression, code, use_calculations, use_calcfunctions, sleep, timeout
     import uuid
     from aiida.orm import Code, Int, Str
     from aiida.engine import run_get_node, submit
-    from lib.expression import generate, validate, evaluate  # pylint: disable=import-error
-    from lib.workchain import generate_outlines, format_outlines, write_workchain  # pylint: disable=import-error
+
+    lib_expression = importlib.import_module('lib.expression')
+    lib_workchain = importlib.import_module('lib.workchain')
 
     if use_calculations and not isinstance(code, Code):
         raise click.BadParameter('if you specify the -C flag, you have to specify a code as well')
 
     if expression is None:
-        expression = generate()
+        expression = lib_expression.generate()
 
-    valid, error = validate(expression)
+    valid, error = lib_expression.validate(expression)
 
     if not valid:
         click.echo(f"the expression '{expression}' is invalid: {error}")
         sys.exit(1)
 
     filename = f'polish_{str(uuid.uuid4().hex)}.py'
-    evaluated = evaluate(expression, modulo)
-    outlines, stack = generate_outlines(expression)
-    outlines_string = format_outlines(outlines, use_calculations, use_calcfunctions)
-    write_workchain(outlines_string, filename=filename)
+    evaluated = lib_expression.evaluate(expression, modulo)
+    outlines, stack = lib_workchain.generate_outlines(expression)
+    outlines_string = lib_workchain.format_outlines(outlines, use_calculations, use_calcfunctions)
+    lib_workchain.write_workchain(outlines_string, filename=filename)
 
     click.echo(f'Expression: {expression}')
 

aiida/cmdline/commands/cmd_restapi.py

@@ -38,7 +38,15 @@
     help='Whether to enable WSGI profiler middleware for finding bottlenecks'
 )
 @click.option('--hookup/--no-hookup', 'hookup', is_flag=True, default=None, help='Hookup app to flask server')
-def restapi(hostname, port, config_dir, debug, wsgi_profile, hookup):
+@click.option(
+    '--posting/--no-posting',
+    'posting',
+    is_flag=True,
+    default=config.CLI_DEFAULTS['POSTING'],
+    help='Enable POST endpoints (currently only /querybuilder).',
+    hidden=True,
+)
+def restapi(hostname, port, config_dir, debug, wsgi_profile, hookup, posting):
     """
     Run the AiiDA REST API server.
 
@@ -55,4 +63,5 @@ def restapi(hostname, port, config_dir, debug, wsgi_profile, hookup):
         debug=debug,
         wsgi_profile=wsgi_profile,
         hookup=hookup,
+        posting=posting,
     )

aiida/restapi/api.py

@@ -25,13 +25,8 @@ class App(Flask):
 
     def __init__(self, *args, **kwargs):
 
-        # Decide whether or not to catch the internal server exceptions (
-        # default is True)
-        catch_internal_server = True
-        try:
-            catch_internal_server = kwargs.pop('catch_internal_server')
-        except KeyError:
-            pass
+        # Decide whether or not to catch the internal server exceptions (default is True)
+        catch_internal_server = kwargs.pop('catch_internal_server', True)
 
         # Basic initialization
         super().__init__(*args, **kwargs)
@@ -95,12 +90,17 @@ def __init__(self, app=None, **kwargs):
           configuration and PREFIX
         """
 
-        from aiida.restapi.resources import ProcessNode, CalcJobNode, Computer, User, Group, Node, ServerInfo
+        from aiida.restapi.common.config import CLI_DEFAULTS
+        from aiida.restapi.resources import (
+            ProcessNode, CalcJobNode, Computer, User, Group, Node, ServerInfo, QueryBuilder
+        )
 
         self.app = app
 
         super().__init__(app=app, prefix=kwargs['PREFIX'], catch_all_404s=True)
 
+        posting = kwargs.pop('posting', CLI_DEFAULTS['POSTING'])
+
         self.add_resource(
             ServerInfo,
             '/',
@@ -111,6 +111,15 @@ def __init__(self, app=None, **kwargs):
             resource_class_kwargs=kwargs
         )
 
+        if posting:
+            self.add_resource(
+                QueryBuilder,
+                '/querybuilder/',
+                endpoint='querybuilder',
+                strict_slashes=False,
+                resource_class_kwargs=kwargs,
+            )
+
         ## Add resources and endpoints to the api
         self.add_resource(
             Computer,

aiida/restapi/common/config.py

@@ -47,4 +47,5 @@
     'WSGI_PROFILE': False,
     'HOOKUP_APP': True,
     'CATCH_INTERNAL_SERVER': False,
+    'POSTING': True,  # Include POST endpoints (currently only /querybuilder)
 }

aiida/restapi/common/identifiers.py

@@ -69,7 +69,7 @@ def construct_full_type(node_type, process_type):
     :return: the full type, which is a unique identifier
     """
     if node_type is None:
-        process_type = ''
+        node_type = ''
 
     if process_type is None:
         process_type = ''

aiida/restapi/resources.py

@@ -207,6 +207,133 @@ def get(self, id=None, page=None):  # pylint: disable=redefined-builtin,invalid-
         return self.utils.build_response(status=200, headers=headers, data=data)
 
 
+class QueryBuilder(BaseResource):
+    """
+    Representation of a QueryBuilder REST API resource (instantiated with a queryhelp JSON).
+
+    It supports POST requests taking in JSON :py:func:`~aiida.orm.querybuilder.QueryBuilder.queryhelp`
+    objects and returning the :py:class:`~aiida.orm.querybuilder.QueryBuilder` result accordingly.
+    """
+    from aiida.restapi.translator.nodes.node import NodeTranslator
+
+    _translator_class = NodeTranslator
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # HTTP Request method decorators
+        if 'get_decorators' in kwargs and isinstance(kwargs['get_decorators'], (tuple, list, set)):
+            self.method_decorators.update({'post': list(kwargs['get_decorators'])})
+
+    def get(self):  # pylint: disable=arguments-differ
+        """Static return to state information about this endpoint."""
+        data = {
+            'message': (
+                'Method Not Allowed. Use HTTP POST requests to use the AiiDA QueryBuilder. '
+                'POST JSON data, which MUST be a valid QueryBuilder.queryhelp dictionary as a JSON object. '
+                'See the documentation at https://aiida.readthedocs.io/projects/aiida-core/en/latest/topics/'
+                'database.html?highlight=QueryBuilder#the-queryhelp for more information.'
+            ),
+        }
+
+        headers = self.utils.build_headers(url=request.url, total_count=1)
+        return self.utils.build_response(
+            status=405,  # Method Not Allowed
+            headers=headers,
+            data={
+                'method': request.method,
+                'url': unquote(request.url),
+                'url_root': unquote(request.url_root),
+                'path': unquote(request.path),
+                'query_string': request.query_string.decode('utf-8'),
+                'resource_type': self.__class__.__name__,
+                'data': data,
+            },
+        )
+
+    def post(self):  # pylint: disable=too-many-branches
+        """
+        POST method to pass query help JSON.
+
+        If the posted JSON is not a valid QueryBuilder queryhelp, the request will fail with an internal server error.
+
+        This uses the NodeTranslator in order to best return Nodes according to the general AiiDA
+        REST API data format, while still allowing the return of other AiiDA entities.
+
+        :return: QueryBuilder result of AiiDA entities in "standard" REST API format.
+        """
+        # pylint: disable=protected-access
+        self.trans._query_help = request.get_json(force=True)
+        # While the data may be correct JSON, it MUST be a single JSON Object,
+        # equivalent of a QuieryBuilder.queryhelp dictionary.
+        assert isinstance(self.trans._query_help, dict), (
+            'POSTed data MUST be a valid QueryBuilder.queryhelp dictionary. '
+            f'Got instead (type: {type(self.trans._query_help)}): {self.trans._query_help}'
+        )
+        self.trans.__label__ = self.trans._result_type = self.trans._query_help['path'][-1]['tag']
+
+        # Handle empty list projections
+        number_projections = len(self.trans._query_help['project'])
+        empty_projections_counter = 0
+        skip_tags = []
+        for tag, projections in tuple(self.trans._query_help['project'].items()):
+            if projections == [{'*': {}}]:
+                self.trans._query_help['project'][tag] = self.trans._default
+            elif not projections:
+                empty_projections_counter += 1
+                skip_tags.append(tag)
+            else:
+                # Use projections as given, no need to "correct" them.
+                pass
+
+        if empty_projections_counter == number_projections:
+            # No projections have been specified in the queryhelp.
+            # To be true to the QueryBuilder response, the last entry in path
+            # is the only entry to be returned, all without edges/links.
+            self.trans._query_help['project'][self.trans.__label__] = self.trans._default
+
+        self.trans.init_qb()
+
+        data = {}
+        if self.trans.get_total_count():
+            if empty_projections_counter == number_projections:
+                # "Normal" REST API retrieval can be used.
+                data = self.trans.get_results()
+            else:
+                # Since the "normal" REST API retrieval relies on single-tag retrieval,
+                # we must instead be more creative with how we retrieve the results here.
+                # So we opt for a dictionary, with the tags being the keys.
+                for tag in self.trans._query_help['project']:
+                    if tag in skip_tags:
+                        continue
+                    self.trans.__label__ = tag
+                    data.update(self.trans.get_formatted_result(tag))
+
+        # Remove 'full_type's when they're `None`
+        for tag, entities in list(data.items()):
+            updated_entities = []
+            for entity in entities:
+                if entity.get('full_type') is None:
+                    entity.pop('full_type', None)
+                updated_entities.append(entity)
+            data[tag] = updated_entities
+
+        headers = self.utils.build_headers(url=request.url, total_count=self.trans.get_total_count())
+        return self.utils.build_response(
+            status=200,
+            headers=headers,
+            data={
+                'method': request.method,
+                'url': unquote(request.url),
+                'url_root': unquote(request.url_root),
+                'path': unquote(request.path),
+                'query_string': request.query_string.decode('utf-8'),
+                'resource_type': self.__class__.__name__,
+                'data': data,
+            },
+        )
+
+
 class Node(BaseResource):
     """
     Differs from BaseResource in trans.set_query() mostly because it takes

aiida/restapi/run_api.py

@@ -39,6 +39,7 @@ def run_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **kwargs)
     :param wsgi_profile: use WSGI profiler middleware for finding bottlenecks in web application
     :param hookup: If true, hook up application to built-in server, else just return it. This parameter
         is deprecated as of AiiDA 1.2.1. If you don't intend to run the API (hookup=False) use `configure_api` instead.
+    :param posting: Whether or not to include POST-enabled endpoints (currently only `/querybuilder`).
 
     :returns: tuple (app, api) if hookup==False or runs app if hookup==True
     """
@@ -80,6 +81,7 @@ def configure_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **k
     :param catch_internal_server:  If true, catch and print internal server errors with full python traceback.
         Useful during app development.
     :param wsgi_profile: use WSGI profiler middleware for finding bottlenecks in the web application
+    :param posting: Whether or not to include POST-enabled endpoints (currently only `/querybuilder`).
 
     :returns: Flask RESTful API
     :rtype: :py:class:`flask_restful.Api`
@@ -89,6 +91,7 @@ def configure_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **k
     config = kwargs.pop('config', CLI_DEFAULTS['CONFIG_DIR'])
     catch_internal_server = kwargs.pop('catch_internal_server', CLI_DEFAULTS['CATCH_INTERNAL_SERVER'])
     wsgi_profile = kwargs.pop('wsgi_profile', CLI_DEFAULTS['WSGI_PROFILE'])
+    posting = kwargs.pop('posting', CLI_DEFAULTS['POSTING'])
 
     if kwargs:
         raise ValueError(f'Unknown keyword arguments: {kwargs}')
@@ -121,4 +124,4 @@ def configure_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **k
         app.wsgi_app = ProfilerMiddleware(app.wsgi_app, restrictions=[30])
 
     # Instantiate and return a Flask RESTful API by associating its app
-    return flask_api(app, **API_CONFIG)
+    return flask_api(app, posting=posting, **config_module.API_CONFIG)

aiida/restapi/translator/nodes/node.py

@@ -563,10 +563,10 @@ def get_formatted_result(self, label):
 
         for node_entry in results[result_name]:
             # construct full_type and add it to every node
-            try:
-                node_entry['full_type'] = construct_full_type(node_entry['node_type'], node_entry['process_type'])
-            except KeyError:
-                node_entry['full_type'] = None
+            node_entry['full_type'] = (
+                construct_full_type(node_entry.get('node_type'), node_entry.get('process_type'))
+                if node_entry.get('node_type') or node_entry.get('process_type') else None
+            )
 
         return results
 

tests/restapi/test_config.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Tests for the configuration options from `aiida.restapi.common.config` when running the REST API."""
+# pylint: disable=redefined-outer-name
+import pytest
+
+
+@pytest.fixture
+def create_app():
+    """Set up Flask App"""
+    from aiida.restapi.run_api import configure_api
+
+    def _create_app(**kwargs):
+        catch_internal_server = kwargs.pop('catch_internal_server', True)
+        api = configure_api(catch_internal_server=catch_internal_server, **kwargs)
+        api.app.config['TESTING'] = True
+        return api.app
+
+    return _create_app
+
+
+def test_posting(create_app):
+    """Test CLI_DEFAULTS['POSTING'] configuration"""
+    from aiida.restapi.common.config import API_CONFIG
+
+    app = create_app(posting=False)
+
+    url = f'{API_CONFIG["PREFIX"]}/querybuilder'
+    for method in ('get', 'post'):
+        with app.test_client() as client:
+            response = getattr(client, method)(url)
+
+        assert response.status_code == 404
+        assert response.status == '404 NOT FOUND'
+
+    del app
+    app = create_app(posting=True)
+
+    url = f'{API_CONFIG["PREFIX"]}/querybuilder'
+    for method in ('get', 'post'):
+        with app.test_client() as client:
+            response = getattr(client, method)(url)
+
+        assert response.status_code != 404
+        assert response.status != '404 NOT FOUND'