add config option to handle multiple dataroots (#1531)

#1513

server/app/app.py

@@ -58,7 +58,7 @@ def cache_control_always(**cache_kwargs):
 
 @webbp.route("/", methods=["GET"])
 @cache_control(public=True, max_age=ONE_WEEK)
-def dataset_index(dataset=None):
+def dataset_index(url_dataroot=None, dataset=None):
     config = current_app.app_config
     if dataset is None:
         if config.single_dataset__datapath:
@@ -66,7 +66,10 @@ def dataset_index(dataset=None):
         else:
             return dataroot_index()
     else:
-        location = path_join(config.multi_dataset__dataroot, dataset)
+        dataroot = config.multi_dataset__dataroot.get(url_dataroot)
+        if dataroot is None:
+            abort(HTTPStatus.NOT_FOUND)
+        location = path_join(dataroot, dataset)
 
     scripts = config.server__scripts
     inline_scripts = config.server__inline_scripts
@@ -91,18 +94,21 @@ def health():
     return health_check(config)
 
 
-def get_data_adaptor(dataset=None):
+def get_data_adaptor(url_dataroot=None, dataset=None):
     config = current_app.app_config
 
     if dataset is None:
         datapath = config.single_dataset__datapath
     else:
-        datapath = path_join(config.multi_dataset__dataroot, dataset)
+        dataroot = config.multi_dataset__dataroot.get(url_dataroot)
+        if dataroot is None:
+            raise DatasetAccessError(f"Invalid dataset {url_dataroot}/{dataset}")
+        datapath = path_join(dataroot, dataset)
         # path_join returns a normalized path.  Therefore it is
         # sufficient to check that the datapath starts with the
         # dataroot to determine that the datapath is under the dataroot.
-        if not datapath.startswith(config.multi_dataset__dataroot):
+        if not datapath.startswith(dataroot):
-            raise DatasetAccessError("Invalid dataset {dataset}")
+            raise DatasetAccessError("Invalid dataset {url_dataroot}/{dataset}")
 
     if datapath is None:
         return common_rest.abort_and_log(HTTPStatus.BAD_REQUEST, "Invalid dataset NONE", loglevel=logging.INFO)
@@ -115,7 +121,7 @@ def rest_get_data_adaptor(func):
     @wraps(func)
     def wrapped_function(self, dataset=None):
         try:
-            with get_data_adaptor(dataset) as data_adaptor:
+            with get_data_adaptor(self.url_dataroot, dataset) as data_adaptor:
                 return func(self, data_adaptor)
         except DatasetAccessError:
             return common_rest.abort_and_log(
@@ -132,22 +138,23 @@ def dataroot_test_index():
     data += "<body><H1>Welcome to cellxgene</H1>"
 
     config = current_app.app_config
-    locator = DataLocator(config.multi_dataset__dataroot, region_name=config.data_locator__s3__region_name)
     datasets = []
-    for fname in locator.ls():
+    for url_dataroot, dataroot in config.multi_dataset__dataroot.items():
-        location = path_join(config.multi_dataset__dataroot, fname)
+        locator = DataLocator(dataroot, region_name=config.data_locator__s3__region_name)
-        try:
+        for fname in locator.ls():
-            MatrixDataLoader(location, app_config=config)
+            location = path_join(dataroot, fname)
-            datasets.append(fname)
+            try:
-        except DatasetAccessError:
+                MatrixDataLoader(location, app_config=config)
-            # skip over invalid datasets
+                datasets.append((url_dataroot, fname))
-            pass
+            except DatasetAccessError:
+                # skip over invalid datasets
+                pass
 
     data += "<br/>Select one of these datasets...<br/>"
     data += "<ul>"
     datasets.sort()
-    for dataset in datasets:
+    for url_dataroot, dataset in datasets:
-        data += f"<li><a href=d/{dataset}>{dataset}</a></li>"
+        data += f"<li><a href={url_dataroot}/{dataset}>{dataset}</a></li>"
     data += "</ul>"
     data += "</body></html>"
 
@@ -165,21 +172,29 @@ def dataroot_index():
         return redirect(config.multi_dataset__index)
 
 
-class SchemaAPI(Resource):
+class DatasetResource(Resource):
+    """Base class for all Resources that act on datasets."""
+
+    def __init__(self, url_dataroot):
+        super().__init__()
+        self.url_dataroot = url_dataroot
+
+
+class SchemaAPI(DatasetResource):
     @cache_control(public=True, max_age=ONE_WEEK)
     @rest_get_data_adaptor
     def get(self, data_adaptor):
         return common_rest.schema_get(data_adaptor, current_app.annotations)
 
 
-class ConfigAPI(Resource):
+class ConfigAPI(DatasetResource):
     @cache_control(public=True, max_age=ONE_WEEK)
     @rest_get_data_adaptor
     def get(self, data_adaptor):
         return common_rest.config_get(current_app.app_config, data_adaptor, current_app.annotations)
 
 
-class AnnotationsObsAPI(Resource):
+class AnnotationsObsAPI(DatasetResource):
     @cache_control(public=True, max_age=ONE_WEEK)
     @rest_get_data_adaptor
     def get(self, data_adaptor):
@@ -191,14 +206,14 @@ class AnnotationsObsAPI(Resource):
         return common_rest.annotations_obs_put(request, data_adaptor, current_app.annotations)
 
 
-class AnnotationsVarAPI(Resource):
+class AnnotationsVarAPI(DatasetResource):
     @cache_control(public=True, max_age=ONE_WEEK)
     @rest_get_data_adaptor
     def get(self, data_adaptor):
         return common_rest.annotations_var_get(request, data_adaptor, current_app.annotations)
 
 
-class DataVarAPI(Resource):
+class DataVarAPI(DatasetResource):
     @cache_control(no_store=True)
     @rest_get_data_adaptor
     def put(self, data_adaptor):
@@ -210,21 +225,21 @@ class DataVarAPI(Resource):
         return common_rest.data_var_get(request, data_adaptor)
 
 
-class ColorsAPI(Resource):
+class ColorsAPI(DatasetResource):
     @cache_control(public=True, max_age=ONE_WEEK)
     @rest_get_data_adaptor
     def get(self, data_adaptor):
         return common_rest.colors_get(data_adaptor)
 
 
-class DiffExpObsAPI(Resource):
+class DiffExpObsAPI(DatasetResource):
     @cache_control(no_store=True)
     @rest_get_data_adaptor
     def post(self, data_adaptor):
         return common_rest.diffexp_obs_post(request, data_adaptor)
 
 
-class LayoutObsAPI(Resource):
+class LayoutObsAPI(DatasetResource):
     @cache_control(public=True, max_age=ONE_WEEK)
     @rest_get_data_adaptor
     def get(self, data_adaptor):
@@ -236,20 +251,25 @@ class LayoutObsAPI(Resource):
         return common_rest.layout_obs_put(request, data_adaptor)
 
 
-def get_api_resources(bp_api):
+def get_api_resources(bp_api, url_dataroot=None):
     api = Api(bp_api)
+
+    def add_resource(resource, url):
+        """convenience function to make the outer function less verbose"""
+        api.add_resource(resource, url, resource_class_args=(url_dataroot,))
+
     # Initialization routes
-    api.add_resource(SchemaAPI, "/schema")
+    add_resource(SchemaAPI, "/schema")
-    api.add_resource(ConfigAPI, "/config")
+    add_resource(ConfigAPI, "/config")
     # Data routes
-    api.add_resource(AnnotationsObsAPI, "/annotations/obs")
+    add_resource(AnnotationsObsAPI, "/annotations/obs")
-    api.add_resource(AnnotationsVarAPI, "/annotations/var")
+    add_resource(AnnotationsVarAPI, "/annotations/var")
-    api.add_resource(DataVarAPI, "/data/var")
+    add_resource(DataVarAPI, "/data/var")
     # Display routes
-    api.add_resource(ColorsAPI, "/colors")
+    add_resource(ColorsAPI, "/colors")
     # Computation routes
-    api.add_resource(DiffExpObsAPI, "/diffexp/obs")
+    add_resource(DiffExpObsAPI, "/diffexp/obs")
-    api.add_resource(LayoutObsAPI, "/layout/obs")
+    add_resource(LayoutObsAPI, "/layout/obs")
     return api
 
 
@@ -285,10 +305,18 @@ class Server:
             # NOTE:  These routes only allow the dataset to be in the directory
             # of the dataroot, and not a subdirectory.  We may want to change
             # the route format at some point
-            bp_api = Blueprint("api_dataset", __name__, url_prefix="/d/<dataset>" + api_version)
+            for url_dataroot in app_config.multi_dataset__dataroot.keys():
-            resources = get_api_resources(bp_api)
+                bp_api = Blueprint(
-            self.app.register_blueprint(resources.blueprint)
+                    f"api_dataset_{url_dataroot}", __name__, url_prefix=f"/{url_dataroot}/<dataset>" + api_version
-            self.app.add_url_rule("/d/<dataset>/", "dataset_index", dataset_index, methods=["GET"])
+                )
+                resources = get_api_resources(bp_api, url_dataroot)
+                self.app.register_blueprint(resources.blueprint)
+                self.app.add_url_rule(
+                    f"/{url_dataroot}/<dataset>/",
+                    f"dataset_index_{url_dataroot}",
+                    lambda dataset: dataset_index(url_dataroot, dataset),
+                    methods=["GET"],
+                )
         self.app.matrix_data_cache_manager = app_config.matrix_data_cache_manager
         self.app.annotations = app_config.user_annotations
         self.app.app_config = app_config

server/common/app_config.py

@@ -1,9 +1,9 @@
 from server import __version__ as cellxgene_version
-from flatten_dict import flatten
+from flatten_dict import flatten, unflatten
 import os
 from os.path import splitext, basename, isdir
 import sys
-from urllib.parse import urlparse
+from urllib.parse import urlparse, quote_plus
 import yaml
 import copy
 
@@ -125,17 +125,23 @@ class AppConfig(object):
         dc = copy.deepcopy(config)
         mapping = {}
 
-        # special case for tiledb_ctx whose value is a dict.
+        # special cases where the value could be a dict.
-        val = config.get("adaptor", {}).get("cxg_adaptor", {}).get("tiledb_ctx")
+        # If its value is not None, the entry is added to the mapping, and not included
-        if val is not None:
+        # in the flattening below.
-            mapping["adaptor__cxg_adaptor__tiledb_ctx"] = (("adaptor", "cxg_adaptor", "tiledb_ctx"), val)
+        dictval_cases = [
-            del dc["adaptor"]["cxg_adaptor"]["tiledb_ctx"]
+            ("adaptor", "cxg_adaptor", "tiledb_ctx"),
-
+            ("server", "csp_directives"),
-        # special case for csp_directives whose value is a dict.
+            ("multi_dataset", "dataroot"),
-        val = config.get("server", {}).get("csp_directives")
+        ]
-        if val is not None:
+        for dictval_case in dictval_cases:
-            mapping["server__csp_directives"] = (("server", "csp_directives"), val)
+            cur = dc
-            del dc["server"]["csp_directives"]
+            for part in dictval_case[:-1]:
+                cur = cur.get(part, {})
+            val = cur.get(dictval_case[-1])
+            if val is not None:
+                key = "__".join(dictval_case)
+                mapping[key] = (dictval_case, val)
+                del cur[dictval_case[-1]]
 
         flat_config = flatten(dc)
         for key, value in flat_config.items():
@@ -162,6 +168,14 @@ class AppConfig(object):
 
         self.is_completed = False
 
+    def write_config(self, config_file):
+        """output the config to a yaml file"""
+        mapping = self.__mapping(self.default_config)
+        for attrname in mapping.keys():
+            mapping[attrname] = getattr(self, attrname)
+        config = unflatten(mapping, splitter=lambda key: key.split("__"))
+        yaml.dump(config, open(config_file, "w"))
+
     def update(self, **kw):
         for key, value in kw.items():
             if not hasattr(self, key):
@@ -302,6 +316,14 @@ class AppConfig(object):
         self.__check_attr("data_locator__s3__region_name", (type(None), bool, str))
         if self.data_locator__s3__region_name is True:
             path = self.single_dataset__datapath or self.multi_dataset__dataroot
+            if type(path) == dict:
+                # if multi_dataset__dataroot is a dict, then use the first key
+                # that is in s3.   NOTE:  it is not supported to have dataroots
+                # in different regions.
+                paths = path.values()
+                for path in paths:
+                    if path.startswith("s3://"):
+                        break
             if path.startswith("s3://"):
                 region_name = discover_s3_region_name(path)
                 if region_name is None:
@@ -366,7 +388,7 @@ class AppConfig(object):
                 )
 
     def handle_multi_dataset(self, context):
-        self.__check_attr("multi_dataset__dataroot", (type(None), str))
+        self.__check_attr("multi_dataset__dataroot", (type(None), dict, str))
         self.__check_attr("multi_dataset__index", (type(None), bool, str))
         self.__check_attr("multi_dataset__allowed_matrix_types", list)
         self.__check_attr("multi_dataset__matrix_cache__max_datasets", int)
@@ -375,6 +397,18 @@ class AppConfig(object):
         if self.multi_dataset__dataroot is None:
             return
 
+        if type(self.multi_dataset__dataroot) == str:
+            self.multi_dataset__dataroot = dict(d=self.multi_dataset__dataroot)
+
+        for key in self.multi_dataset__dataroot.keys():
+            # sanity check for well formed keys
+            if type(key) != str:
+                raise ConfigurationError(f"error in multi_dataset__dataroot {key}")
+            if quote_plus(key) != key:
+                raise ConfigurationError(f"error in multi_dataset__dataroot {key}")
+            if os.path.split(os.path.normpath(key))[-1] != key:
+                raise ConfigurationError(f"error in multi_dataset__dataroot {key}")
+
         # error checking
         for mtype in self.multi_dataset__allowed_matrix_types:
             try:

server/common/default_config.py

@@ -29,6 +29,23 @@ presentation:
   custom_colors: true
 
 multi_dataset:
+  # If dataroot is set, then cellxgene may serve multiple datasets.  This parameter is not
+  # compatable with single_dataset/datapath.
+  # dataroot may be a string, representing the path to a directory or S3 prefix.  In this
+  # case the datasets in that location are accessed from <server>/d/<datasetname>.
+  # example:
+  #     dataroot:  /path/to/datasets/
+  # or
+  #     dataroot:  s3://bucket/prefix/
+  #
+  # As an alternative, dataroot can be a dictionary, mapping url prefixes to dataroot paths.
+  # example:
+  #     dataroot:
+  #         set1 : /path/to/set1_datasets/
+  #         set2 : /path/to/set2_datasets/
+  # In this case, datasets can be accessed from <server>/set1/<datasetname> or
+  # <server>/set2/<datasetname>.
+
   dataroot: null
 
   # The index page when in multi-dataset mode:

server/common/health.py

@@ -23,12 +23,13 @@ def health_check(config):
     """
     health = {"status": None, "version": "1", "releaseID": cellxgene_version}
 
-    checks = [
+    checks = False
-        (config.single_dataset__datapath is not None or config.multi_dataset__dataroot is not None),
+    if config.single_dataset__datapath is not None:
-        _is_accessible(config.single_dataset__datapath, config),
+        checks = _is_accessible(config.single_dataset__datapath, config)
-        _is_accessible(config.multi_dataset__dataroot, config),
+    elif config.multi_dataset__dataroot is not None:
-    ]
+        checks = all([_is_accessible(datapath, config) for datapath in config.multi_dataset__dataroot.values()])
-    health["status"] = "pass" if all(checks) else "fail"
+
+    health["status"] = "pass" if checks else "fail"
     code = HTTPStatus.OK if health["status"] == "pass" else HTTPStatus.BAD_REQUEST
     response = make_response(jsonify(health), code)
     response.headers["Content-Type"] = "application/health+json"

server/test/__init__.py

@@ -2,13 +2,19 @@ import random
 import shutil
 import string
 import tempfile
+import requests
+import time
+import os
+from subprocess import Popen
 from os import path, popen
+from contextlib import contextmanager
 
 import pandas as pd
 
 from server.common.annotations import AnnotationsLocalFile
 from server.common.data_locator import DataLocator
-from server.common.app_config import AppConfig
+from server.common.app_config import AppConfig, DEFAULT_SERVER_PORT
+from server.common.utils import find_available_port
 from server.data_common.fbs.matrix import encode_matrix_fbs
 from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataType
 
@@ -81,3 +87,54 @@ def app_config(data_locator, backed=False, extra={}):
 
 def random_string(n):
     return "".join(random.choice(string.ascii_letters) for _ in range(n))
+
+
+@contextmanager
+def test_server(command_line_args=[], app_config=None):
+    """A context to run the cellxgene server.
+    Command line arguments can be passed in, as well as an app_config.
+    This function is meant to be used like this, for example:
+
+    with test_server(...) as server:
+       r = requests.get(f"{server}/...")
+       // check r
+
+    where the server can be accessed within the context, and is terminated when
+    the context is exited.
+    The port is automatically set using find_available_port.
+    The verbose flag is automatically set to True.
+    If an app_config is provided, then this function writes a temporary
+    yaml config file, which this server will read and parse.
+    """
+
+    port = DEFAULT_SERVER_PORT
+    port = find_available_port("localhost", port)
+    command = ["cellxgene", "--no-upgrade-check", "launch", "--verbose", "--port=%d" % port] + command_line_args
+
+    tempdir = None
+    if app_config:
+        tempdir = tempfile.TemporaryDirectory()
+        config_file = os.path.join(tempdir.name, "config.yaml")
+        app_config.write_config(config_file)
+        command.extend(["-c", config_file])
+
+    server = f"http://localhost:{port}"
+    ps = Popen(command)
+
+    for _ in range(10):
+        try:
+            requests.get(f"{server}/health")
+            break
+        except requests.exceptions.ConnectionError:
+            time.sleep(1)
+
+    if tempdir:
+        tempdir.cleanup()
+
+    try:
+        yield server
+    finally:
+        try:
+            ps.terminate()
+        except ProcessLookupError:
+            pass

server/test/test_app_config.py

@@ -1,5 +1,8 @@
 import unittest
 from server.common.app_config import AppConfig
+from server.common.errors import ConfigurationError
+from server.test import PROJECT_ROOT, test_server
+import requests
 
 # NOTE, there are more tests that should be written for AppConfig.
 # this is just a start.
@@ -26,3 +29,43 @@ class AppConfigTest(unittest.TestCase):
         c.update(server__scripts=("a", "b"), server__inline_scripts=["c", "d"])
         v = c.changes_from_default()
         self.assertCountEqual(v, [("server__scripts", ["a", "b"], []), ("server__inline_scripts", ["c", "d"], [])])
+
+    def test_multi_dataset(self):
+
+        c = AppConfig()
+        # test for illegal url_dataroots
+        for illegal in ("a/b", "../b", "!$*", "\\n", "", "(bad)"):
+            c.update(multi_dataset__dataroot={illegal: f"{PROJECT_ROOT}/example-dataset"})
+            with self.assertRaises(ConfigurationError):
+                c.complete_config()
+
+        # test for legal url_dataroots
+        for legal in (
+            "d",
+            "this.is-okay_",
+        ):
+            c.update(multi_dataset__dataroot={legal: f"{PROJECT_ROOT}/example-dataset"})
+            c.complete_config()
+
+        # test that multi dataroots work end to end
+        c.update(
+            multi_dataset__dataroot=dict(
+                set1=f"{PROJECT_ROOT}/example-dataset",
+                set2=f"{PROJECT_ROOT}/server/test/test_datasets"
+            )
+        )
+        c.complete_config()
+
+        with test_server(app_config=c) as server:
+            session = requests.Session()
+
+            r = session.get(f"{server}/set1/pbmc3k.h5ad/api/v0.2/config")
+            data_config = r.json()
+            assert data_config["config"]["displayNames"]["dataset"] == "pbmc3k"
+
+            r = session.get(f"{server}/set2/pbmc3k.cxg/api/v0.2/config")
+            data_config = r.json()
+            assert data_config["config"]["displayNames"]["dataset"] == "pbmc3k"
+
+            r = session.get(f"{server}/health")
+            assert r.json()["status"] == "pass"