Create hosted user annotations [1685] (#1726)

* add function to retrieve latest annotation from db, db updates

* read and write tiledb arrays

* adding tests

server/common/annotations/annotations.py

@@ -0,0 +1,78 @@
+from abc import ABCMeta, abstractmethod
+
+import fastobo
+import fsspec
+from server.common.errors import OntologyLoadFailure
+
+from server.common.utils import series_to_schema
+
+
+class Annotations(metaclass=ABCMeta):
+    """ baseclass for annotations, including ontologies"""
+
+    """ our default ontology is the PURL for the Cell Ontology.
+    See http://www.obofoundry.org/ontology/cl.html """
+    DefaultOnotology = "http://purl.obolibrary.org/obo/cl.obo"
+
+    def __init__(self):
+        self.ontology_data = None
+
+    def load_ontology(self, path):
+        """Load and parse ontologies - currently support OBO files only."""
+        if path is None:
+            path = self.DefaultOnotology
+
+        try:
+            with fsspec.open(path) as f:
+                obo = fastobo.iter(f)
+                terms = filter(lambda stanza: type(stanza) is fastobo.term.TermFrame, obo)
+                names = [tag.name for term in terms for tag in term if type(tag) is fastobo.term.NameClause]
+                self.ontology_data = names
+
+        except FileNotFoundError as e:
+            raise OntologyLoadFailure("Unable to find OBO ontology path") from e
+
+        except SyntaxError as e:
+            raise OntologyLoadFailure("Syntax error loading OBO ontology") from e
+
+        except Exception as e:
+            raise OntologyLoadFailure("Error loading OBO file") from e
+
+    def get_schema(self, data_adaptor):
+        schema = []
+        labels = self.read_labels(data_adaptor)
+        if labels is not None and not labels.empty:
+            for col in labels.columns:
+                col_schema = dict(name=col, writable=True)
+                col_schema.update(series_to_schema(labels[col]))
+                schema.append(col_schema)
+
+        return schema
+
+    @abstractmethod
+    def set_collection(self, name):
+        """set or create a new annotation collection"""
+        pass
+
+    @abstractmethod
+    def read_labels(self, data_adaptor):
+        """Return the labels as a pandas.DataFrame"""
+        pass
+
+    @abstractmethod
+    def write_labels(self, df, data_adaptor):
+        """Write the labels (df) to a persistent storage such that it can later be read"""
+        pass
+
+    def update_parameters(self, parameters, data_adaptor):
+        """Update configuration parameters that describe information about the annotations feature"""
+        params = {}
+        params["annotations"] = True
+
+        if self.ontology_data:
+            params["annotations_cell_ontology_enabled"] = True
+            params["annotations_cell_ontology_terms"] = self.ontology_data
+        else:
+            params["annotations_cell_ontology_enabled"] = False
+
+        parameters.update(params)

server/common/annotations/hosted_tiledb.py

@@ -0,0 +1,113 @@
+import json
+import os
+import re
+import time
+
+import pandas as pd
+import tiledb
+from flask import current_app
+
+from server.common.annotations.annotations import Annotations
+from server.converters.cxgtool import sanitize_keys, generate_schema_hints_and_convert_value_types, cxg_dtype
+from server.db.cellxgene_orm import CellxGeneDataset, Annotation
+
+
+class AnnotationsHostedTileDB(Annotations):
+    CXG_ANNO_COLLECTION = "cxg_anno_collection"
+
+    def __init__(self, directory_path, db):
+        super().__init__()
+        self.db = db
+        self.directory_path = directory_path
+
+    def check_category_names(self, df):
+        sanitize_keys(df.keys().to_list(), False)
+
+    def is_safe_collection_name(self, name):
+        """
+        return true if this is a safe collection name
+        this is ultra conservative. If we want to allow full legal file name syntax,
+        we could look at modules like `pathvalidate`
+        """
+        if name is None:
+            return False
+        return re.match(r"^[\w\-]+$", name) is not None
+
+    def set_collection(self, name):
+        self.CXG_ANNO_COLLECTION = name
+
+    def read_labels(self, data_adaptor):
+        user_id = current_app.auth.get_user_id()
+        dataset_name = data_adaptor.get_location()
+        dataset_id = str(self.db.query(
+            table_args=[CellxGeneDataset],
+            filter_args=[CellxGeneDataset.name == dataset_name]
+        )[0].id)
+
+        annotation_object = self.db.query_for_most_recent(
+            Annotation, [Annotation.user_id == user_id, Annotation.dataset_id == dataset_id]
+        )
+        if annotation_object:
+            df = tiledb.open(annotation_object.tiledb_uri)
+            pandas_df = self.convert_to_pandas_df(df)
+            return pandas_df
+        else:
+            return None
+
+    def convert_to_pandas_df(self, tileDBArray):
+        repr_meta = None
+        index_dims = None
+        if '__pandas_attribute_repr' in tileDBArray.meta:
+            # backwards compatibility... unsure if necessary at this point
+            repr_meta = json.loads(tileDBArray.meta['__pandas_attribute_repr'])
+        if '__pandas_index_dims' in tileDBArray.meta:
+            index_dims = json.loads(tileDBArray.meta['__pandas_index_dims'])
+
+        data = tileDBArray[:]
+        indexes = list()
+
+        for col_name, col_val in data.items():
+            if repr_meta and col_name in repr_meta:
+                new_col = pd.Series(col_val, dtype=repr_meta[col_name])
+                data[col_name] = new_col
+            elif index_dims and col_name in index_dims:
+                new_col = pd.Series(col_val, dtype=index_dims[col_name])
+                data[col_name] = new_col
+                indexes.append(col_name)
+
+        new_df = pd.DataFrame.from_dict(data)
+        if len(indexes) > 0:
+            new_df.set_index(indexes, inplace=True)
+
+        return new_df
+
+    def write_labels(self, df, data_adaptor):
+
+        user_id = current_app.auth.get_user_id()
+        timestamp = time.time()
+        dataset_name = data_adaptor.get_location()
+        dataset_id = self.db.get_or_create_dataset(dataset_name)
+        user_id = self.db.get_or_create_user(user_id)
+
+        uri = f"{self.directory_path}-{dataset_name}-{user_id}-{timestamp}"
+        if uri.startswith("s3://"):
+            pass
+        else:
+            os.makedirs(uri, exist_ok=True)
+        schema_hints, values = generate_schema_hints_and_convert_value_types(df)
+
+        annotation = Annotation(
+            tiledb_uri=uri,
+            user_id=user_id,
+            dataset_id=str(dataset_id),
+            schema_hints=json.dumps(schema_hints)
+        )
+        if not df.empty:
+            self.check_category_names(df)
+            # convert to tiledb datatypes
+            for col in df:
+                df[col] = df[col].astype(cxg_dtype(df[col]))
+            tiledb.from_pandas(uri, df)
+
+        self.db.session.add(annotation)
+        self.db.session.commit()

server/common/annotations/local_file_csv.py

@@ -1,86 +1,16 @@
-import json
-import uuid
-import time
-from datetime import datetime
-import re
-import os
-import pandas as pd
-from hashlib import blake2b
 import base64
-from server import __version__ as cellxgene_version
+import os
+import re
 import threading
-from server.common.errors import AnnotationsError, OntologyLoadFailure
-from server.common.utils import series_to_schema
-import fsspec
-import fastobo
-from flask import session, current_app, has_request_context
-from abc import ABCMeta, abstractmethod
+from datetime import datetime
+from hashlib import blake2b
 
-from server.db.cellxgene_orm import CellxGeneDataset, Annotation
-from server.db.db_utils import DbUtils
+import pandas as pd
+from flask import session, has_request_context, current_app
 
-
-class Annotations(metaclass=ABCMeta):
-    """ baseclass for annotations, including ontologies"""
-
-    """ our default ontology is the PURL for the Cell Ontology.
-    See http://www.obofoundry.org/ontology/cl.html """
-    DefaultOnotology = "http://purl.obolibrary.org/obo/cl.obo"
-
-    def __init__(self):
-        self.ontology_data = None
-
-    def load_ontology(self, path):
-        """Load and parse ontologies - currently support OBO files only."""
-        if path is None:
-            path = self.DefaultOnotology
-
-        try:
-            with fsspec.open(path) as f:
-                obo = fastobo.iter(f)
-                terms = filter(lambda stanza: type(stanza) is fastobo.term.TermFrame, obo)
-                names = [tag.name for term in terms for tag in term if type(tag) is fastobo.term.NameClause]
-                self.ontology_data = names
-
-        except FileNotFoundError as e:
-            raise OntologyLoadFailure("Unable to find OBO ontology path") from e
-
-        except SyntaxError as e:
-            raise OntologyLoadFailure("Syntax error loading OBO ontology") from e
-
-        except Exception as e:
-            raise OntologyLoadFailure("Error loading OBO file") from e
-
-    def get_schema(self, data_adaptor):
-        schema = []
-        labels = self.read_labels(data_adaptor)
-        if labels is not None and not labels.empty:
-            for col in labels.columns:
-                col_schema = dict(name=col, writable=True)
-                col_schema.update(series_to_schema(labels[col]))
-                schema.append(col_schema)
-
-        return schema
-
-    @abstractmethod
-    def set_collection(self, name):
-        """set or create a new annotation collection"""
-        pass
-
-    @abstractmethod
-    def read_labels(self, data_adaptor):
-        """Return the labels as a pandas.DataFrame"""
-        pass
-
-    @abstractmethod
-    def write_labels(self, df, data_adaptor):
-        """Write the labels (df) to a persistent storage such that it can later be read"""
-        pass
-
-    @abstractmethod
-    def update_parameters(self, parameters, data_adaptor):
-        """Update configuration parameters that describe information about the annotations feature"""
-        pass
+from server import __version__ as cellxgene_version
+from server.common.annotations.annotations import Annotations
+from server.common.errors import AnnotationsError
 
 
 class AnnotationsLocalFile(Annotations):
@@ -101,7 +31,6 @@ class AnnotationsLocalFile(Annotations):
     def is_safe_collection_name(self, name):
         """
         return true if this is a safe collection name
-
         this is ultra conservative. If we want to allow full legal file name syntax,
         we could look at modules like `pathvalidate`
         """
@@ -265,55 +194,3 @@ class AnnotationsLocalFile(Annotations):
                 params["annotations-data-collection-name"] = collection
 
         parameters.update(params)
-
-
-class AnnotationsHostedTileDB(Annotations):
-    def __init__(self, directory_path: str, db: DbUtils):
-        super().__init__()
-        self.db = db
-        self.directory_path = directory_path
-
-    def set_collection(self, name):
-        pass
-
-    def read_labels(self, data_adaptor):
-        uid = current_app.auth.get_user_id()
-        dataset_name = data_adaptor.get_location()
-        dataset = self.db.query(table_args=[CellxGeneDataset], filter_args=[CellxGeneDataset.name == dataset_name])
-        # Todo @madison retrieve latest based on timestamp
-        annotation_object = self.db.query_for_most_recent(  # noqa F841
-            Annotation, [Annotation.user_id == uid, Annotation.dataset == dataset]
-        )
-        # Todo in future pr, retrieve dataframe from tiledb uri
-
-    def write_labels(self, df, data_adaptor):
-        uid = current_app.auth.get_user_id()
-        timestamp = time.time()
-        dataset_name = data_adaptor.get_location()
-        try:
-            dataset_id = self.db.query(
-                table_args=[CellxGeneDataset], filter_args=[CellxGeneDataset.name == dataset_name]
-            )[0].id
-        except IndexError:
-            dataset_id = uuid.uuid4()
-            dataset = CellxGeneDataset(id=dataset_id, name=dataset_name)
-            self.db.session.add(dataset)
-
-        uri = f"{self.directory_path}/{dataset_name}/{uid}/{timestamp}"
-        if "s3" in uri:
-            pass
-        else:
-            os.makedirs(uri, exist_ok=True)
-        schema_hints = {}
-        annotation = Annotation(
-            tiledb_uri=uri,
-            user_id=uid,
-            dataset_id=str(dataset_id),
-            schema_hints=json.dumps(schema_hints)
-        )
-        # todo in future pr -- write df to tiledb, store at uri
-        self.db.session.add(annotation)
-        self.db.session.commit()
-
-    def update_parameters(self, parameters, data_adaptor):
-        pass

server/common/app_config.py

@@ -12,11 +12,13 @@ from server.common.errors import ConfigurationError, DatasetAccessError, Ontolog
 from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataCacheManager, MatrixDataType
 from server.common.utils import find_available_port, is_port_available
 import warnings
-from server.common.annotations import AnnotationsLocalFile
+from server.common.annotations.hosted_tiledb import AnnotationsHostedTileDB
+from server.common.annotations.local_file_csv import AnnotationsLocalFile
 from server.common.utils import custom_format_warning
 import server.compute.diffexp_cxg as diffexp_tiledb
 from server.common.data_locator import discover_s3_region_name
 from server.auth.auth import AuthTypeFactory
+from server.db.db_utils import DbUtils
 
 DEFAULT_SERVER_PORT = 5005
 # anything bigger than this will generate a special message
@@ -277,6 +279,7 @@ class AppConfig(object):
                 "is_authenticated": auth.is_user_authenticated(),
                 "requires_client_login": auth.requires_client_login(),
                 "username": auth.get_user_name(),
+                "user_id": auth.get_user_id()
             }
             if auth.requires_client_login():
                 config["authentication"].update({
@@ -736,6 +739,8 @@ class DatasetConfig(BaseConfig):
             self.user_annotations__local_file_csv__file = dc["user_annotations"]["local_file_csv"]["file"]
             self.user_annotations__ontology__enable = dc["user_annotations"]["ontology"]["enable"]
             self.user_annotations__ontology__obo_location = dc["user_annotations"]["ontology"]["obo_location"]
+            self.user_annotations__hosted_tiledb_array__db_uri = dc["user_annotations"]["hosted_tiledb_array"]["db_uri"]
+            self.user_annotations__hosted_tiledb_array__hosted_file_directory = dc["user_annotations"]["hosted_tiledb_array"]["hosted_file_directory"]  # noqa E501
 
             self.embeddings__names = dc["embeddings"]["names"]
             self.embeddings__enable_reembedding = dc["embeddings"]["enable_reembedding"]
@@ -786,6 +791,8 @@ class DatasetConfig(BaseConfig):
         self.check_attr("user_annotations__local_file_csv__file", (type(None), str))
         self.check_attr("user_annotations__ontology__enable", bool)
         self.check_attr("user_annotations__ontology__obo_location", (type(None), str))
+        self.check_attr("user_annotations__hosted_tiledb_array__db_uri", (type(None), str))
+        self.check_attr("user_annotations__hosted_tiledb_array__hosted_file_directory", (type(None), str))
 
         if self.user_annotations__enable:
             server_config = self.app_config.server_config
@@ -797,43 +804,49 @@ class DatasetConfig(BaseConfig):
 
             # TODO, replace this with a factory pattern once we have more than one way
             # to do annotations.  currently only local_file_csv
-            if self.user_annotations__type != "local_file_csv":
-                raise ConfigurationError('The only annotation type support is "local_file_csv"')
+            if self.user_annotations__type == "local_file_csv":
+                dirname = self.user_annotations__local_file_csv__directory
+                filename = self.user_annotations__local_file_csv__file
 
-            dirname = self.user_annotations__local_file_csv__directory
-            filename = self.user_annotations__local_file_csv__file
+                if filename is not None and dirname is not None:
+                    raise ConfigurationError("'annotations-file' and 'annotations-dir' may not be used together.")
 
-            if filename is not None and dirname is not None:
-                raise ConfigurationError("'annotations-file' and 'annotations-dir' may not be used together.")
+                if filename is not None:
+                    lf_name, lf_ext = splitext(filename)
+                    if lf_ext and lf_ext != ".csv":
+                        raise ConfigurationError(f"annotation file type must be .csv: {filename}")
 
-            if filename is not None:
-                lf_name, lf_ext = splitext(filename)
-                if lf_ext and lf_ext != ".csv":
-                    raise ConfigurationError(f"annotation file type must be .csv: {filename}")
+                if dirname is not None and not isdir(dirname):
+                    try:
+                        os.mkdir(dirname)
+                    except OSError:
+                        raise ConfigurationError("Unable to create directory specified by --annotations-dir")
 
-            if dirname is not None and not isdir(dirname):
-                try:
-                    os.mkdir(dirname)
-                except OSError:
-                    raise ConfigurationError("Unable to create directory specified by --annotations-dir")
+                self.user_annotations = AnnotationsLocalFile(dirname, filename)
 
-            self.user_annotations = AnnotationsLocalFile(dirname, filename)
-
-            # if the user has specified a fixed label file, go ahead and validate it
-            # so that we can remove errors early in the process.
-            server_config = self.app_config.server_config
-            if server_config.single_dataset__datapath and self.user_annotations__local_file_csv__file:
-                with server_config.matrix_data_cache_manager.data_adaptor(
-                    self.tag, server_config.single_dataset__datapath, self.app_config
-                ) as data_adaptor:
-                    data_adaptor.check_new_labels(self.user_annotations.read_labels(data_adaptor))
-
-            if self.user_annotations__ontology__enable or self.user_annotations__ontology__obo_location:
-                try:
-                    self.user_annotations.load_ontology(self.user_annotations__ontology__obo_location)
-                except OntologyLoadFailure as e:
-                    raise ConfigurationError("Unable to load ontology terms\n" + str(e))
+                # if the user has specified a fixed label file, go ahead and validate it
+                # so that we can remove errors early in the process.
+                server_config = self.app_config.server_config
+                if server_config.single_dataset__datapath and self.user_annotations__local_file_csv__file:
+                    with server_config.matrix_data_cache_manager.data_adaptor(
+                            self.tag, server_config.single_dataset__datapath, self.app_config
+                    ) as data_adaptor:
+                        data_adaptor.check_new_labels(self.user_annotations.read_labels(data_adaptor))
 
+                if self.user_annotations__ontology__enable or self.user_annotations__ontology__obo_location:
+                    try:
+                        self.user_annotations.load_ontology(self.user_annotations__ontology__obo_location)
+                    except OntologyLoadFailure as e:
+                        raise ConfigurationError("Unable to load ontology terms\n" + str(e))
+            elif self.user_annotations__type == "hosted_tiledb_array":
+                self.check_attr("user_annotations__hosted_tiledb_array__db_uri", str)
+                self.check_attr("user_annotations__hosted_tiledb_array__hosted_file_directory", str)
+                self.user_annotations = AnnotationsHostedTileDB(
+                    directory_path=self.user_annotations__hosted_tiledb_array__hosted_file_directory,
+                    db=DbUtils(self.user_annotations__hosted_tiledb_array__db_uri),
+                )
+            else:
+                raise ConfigurationError('The only annotation type support is "local_file_csv" or "hosted_tiledb_array')
         else:
             if self.user_annotations__type == "local_file_csv":
                 dirname = self.user_annotations__local_file_csv__directory

server/common/default_config.py

@@ -171,6 +171,9 @@ dataset:
   user_annotations:
     enable: true
     type: local_file_csv
+    hosted_tiledb_array:
+        db_uri: null
+        hosted_file_directory: null
     local_file_csv:
       directory: null
       file: null

server/common/errors.py

@@ -46,6 +46,11 @@ define_request_exception(
     "Raised when there is an authentication error",
     default_status_code=HTTPStatus.UNAUTHORIZED)
 
+define_request_exception(
+    "AnnotationCategoryNameError",
+    "Raised when an annotation category name cant be saved",
+    default_status_code=HTTPStatus.UNPROCESSABLE_ENTITY)
+
 define_exception("OntologyLoadFailure", "Raised when reading the ontology file fails")
 define_exception("ConfigurationError", "Raised when checking configuration errors")
 define_exception("PrepareError", "Raised when data is misprepared")

server/converters/cxgtool.py

@@ -15,7 +15,7 @@ import json
 from scipy.stats import mode
 
 from server.common.colors import convert_anndata_category_colors_to_cxg_category_colors
-from server.common.errors import ColorFormatException
+from server.common.errors import ColorFormatException, AnnotationCategoryNameError
 from server.common.corpora import (
     corpora_get_props_from_anndata,
     corpora_get_versions_from_anndata,
@@ -291,20 +291,26 @@ def alias_index_col(df, df_name, index_col_name):
     return (df, index_col_name)
 
 
+def generate_schema_hints_and_convert_value_types(df):
+    value = {}
+    schema_hints = {}
+    for k, v in df.items():
+        dtype, hints = cxg_type(v)
+        value[k] = v.to_numpy(dtype=dtype)
+        if hints:
+            schema_hints.update({k: hints})
+    return schema_hints, value
+
+
 def save_dataframe(container, name, df, index_col_name, ctx):
     A_name = f"{container}/{name}"
     (df, index_col_name) = alias_index_col(df, name, index_col_name)
     create_dataframe(A_name, df, ctx=ctx)
     with tiledb.DenseArray(A_name, mode="w", ctx=ctx) as A:
-        value = {}
-        schema_hints = {}
-        for k, v in df.items():
-            dtype, hints = cxg_type(v)
-            value[k] = v.to_numpy(dtype=dtype)
-            if hints:
-                schema_hints.update({k: hints})
-
+        schema_hints, value = generate_schema_hints_and_convert_value_types(df)
         schema_hints.update({"index": index_col_name})
+        # convert all values in all cols to a numpy version of cxg datatypes,
+        # then store the contents in the tiledb array A
         A[:] = value
         A.meta["cxg_schema"] = json.dumps(schema_hints)
 
@@ -598,7 +604,7 @@ def create_cxg_group_metadata(adata, basefname, title=None, about=None, corpora_
     return cxg_group_metadata
 
 
-def sanitize_keys(keys):
+def sanitize_keys(keys, update_keys=True):
     """
     We need names to be safe to use as attribute names in tiledb.  See:
         TileDB-Inc/TileDB#1575
@@ -635,6 +641,8 @@ def sanitize_keys(keys):
 
     for k, v, in clean_unique_keys.items():
         if k != v:
+            if update_keys is False:
+                raise AnnotationCategoryNameError(f"{k} not a valid category name, please resubmit")
             log(1, f"Renaming {k} to {v}")
     return clean_unique_keys
 

server/db/db_utils.py

@@ -1,9 +1,10 @@
 import typing
+import uuid
 
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
-from server.db.cellxgene_orm import Base
+from server.db.cellxgene_orm import Base, CellxGeneDataset, CellxGeneUser
 
 
 class DbUtils:
@@ -34,7 +35,33 @@ class DbUtils:
         )
 
     def query_for_most_recent(self, table: Base, filter_args: typing.List[bool] = None) -> Base:
-        return self.session.query(table).filter(*filter_args).order_by(table.created_at.desc()).limit(1).all()[0]
+        try:
+            return self.session.query(table).filter(*filter_args).order_by(table.created_at.desc()).limit(1).all()[0]
+        except IndexError:
+            return None
+
+    def get_or_create_dataset(self, dataset_name):
+        try:
+            dataset_id = self.query(
+                table_args=[CellxGeneDataset], filter_args=[CellxGeneDataset.name == dataset_name]
+            )[0].id
+        except IndexError:
+            dataset_id = uuid.uuid4()
+            dataset = CellxGeneDataset(id=dataset_id, name=dataset_name)
+            self.session.add(dataset)
+            self.session.commit()
+        return str(dataset_id)
+
+    def get_or_create_user(self, user_id):
+        try:
+            user_id = self.query(
+                table_args=[CellxGeneUser], filter_args=[CellxGeneUser.id == user_id]
+            )[0].id
+        except IndexError:
+            user = CellxGeneUser(id=user_id)
+            self.session.add(user)
+            self.session.commit()
+        return str(user_id)
 
 
 class DBSessionMaker:

server/eb/app.py

@@ -73,7 +73,8 @@ def handle_config_from_secret(app_config):
 
     keyattrs = (
         ("flask_secret_key", "app__flask_secret_key"),
-        ("oauth_client_secret", "authentication__params_oauth__client_secret")
+        ("oauth_client_secret", "authentication__params_oauth__client_secret"),
+        ("db_uri", "user_annotations__hosted_tiledb_array__db_uri"),
     )
 
     for key, attr in keyattrs:

server/test/__init__.py

@@ -11,18 +11,50 @@ from contextlib import contextmanager
 
 import pandas as pd
 
-from server.common.annotations import AnnotationsLocalFile
+from server.common.annotations.hosted_tiledb import AnnotationsHostedTileDB
+from server.common.annotations.local_file_csv import AnnotationsLocalFile
+
 from server.common.data_locator import DataLocator
 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
-
+from server.db.db_utils import DbUtils
 
 PROJECT_ROOT = popen("git rev-parse --show-toplevel").read().strip()
 FIXTURES_ROOT = PROJECT_ROOT + "/server/test/fixtures"
 
 
+def data_with_tmp_tiledb_annotations(ext: MatrixDataType):
+    tmp_dir = tempfile.mkdtemp()
+    fname = {
+        MatrixDataType.H5AD: f"{PROJECT_ROOT}/example-dataset/pbmc3k.h5ad",
+        MatrixDataType.CXG: "test/fixtures/pbmc3k.cxg",
+    }[ext]
+    data_locator = DataLocator(fname)
+    config = AppConfig()
+    config.update_server_config(
+        multi_dataset__dataroot=data_locator.path, authentication__type="test"
+    )
+    config.update_default_dataset_config(
+        embeddings__names=["umap"],
+        presentation__max_categories=100,
+        diffexp__lfc_cutoff=0.01,
+        user_annotations__type="hosted_tiledb_array",
+        user_annotations__hosted_tiledb_array__db_uri="postgresql://postgres:test_pw@localhost:5432",
+        user_annotations__hosted_tiledb_array__hosted_file_directory=tmp_dir
+    )
+
+    config.complete_config()
+
+    data = MatrixDataLoader(data_locator.abspath()).open(config)
+    annotations = AnnotationsHostedTileDB(
+        tmp_dir,
+        DbUtils("postgresql://postgres:test_pw@localhost:5432")
+    )
+    return data, tmp_dir, annotations
+
+
 def data_with_tmp_annotations(ext: MatrixDataType, annotations_fixture=False):
     tmp_dir = tempfile.mkdtemp()
     annotations_file = path.join(tmp_dir, "test_annotations.csv")
@@ -40,6 +72,7 @@ def data_with_tmp_annotations(ext: MatrixDataType, annotations_fixture=False):
     config.update_default_dataset_config(
         embeddings__names=["umap"], presentation__max_categories=100, diffexp__lfc_cutoff=0.01,
     )
+
     config.complete_config()
     data = MatrixDataLoader(data_locator.abspath()).open(config)
     annotations = AnnotationsLocalFile(None, annotations_file)

server/test/unit/common/test_writable_annotation.py

@@ -1,6 +1,11 @@
 import json
 from os import path, listdir
 import unittest
+from unittest.mock import MagicMock, patch
+
+import tiledb
+from flask import Flask
+
 import server.test.unit.decode_fbs as decode_fbs
 import shutil
 
@@ -8,8 +13,134 @@ import numpy as np
 import pandas as pd
 
 from server.common.rest import schema_get_helper, annotations_put_fbs_helper
-from server.test import data_with_tmp_annotations, make_fbs
+from server.db.cellxgene_orm import CellxGeneDataset, Annotation
+from server.test import data_with_tmp_annotations, make_fbs, data_with_tmp_tiledb_annotations
 from server.data_common.matrix_loader import MatrixDataType
+from server.common.errors import AnnotationCategoryNameError
+
+
+class auth(object):
+    def get_user_id():
+        return "1234"
+
+
+class WritableTileDBStoredAnnotationTest(unittest.TestCase):
+    def setUp(self):
+        self.user_id = '1234'
+        self.data, self.tmp_dir, self.annotations = data_with_tmp_tiledb_annotations(MatrixDataType.H5AD)
+        self.data.dataset_config.user_annotations = self.annotations
+        self.db = self.annotations.db
+        self.n_rows = self.data.get_shape()[0]
+        self.test_dict = {
+            "cat_A": pd.Series(["label_A"] * self.n_rows, dtype="category"),
+            "cat_B": pd.Series(["label_B"] * self.n_rows, dtype="category"),
+        }
+        self.fbs = make_fbs(self.test_dict)
+        self.df = pd.DataFrame(self.test_dict)
+        self.app = Flask('fake_app')
+        self.app.__setattr__("auth", auth)
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def annotation_put_fbs(self, fbs):
+        annotations_put_fbs_helper(self.data, fbs)
+        res = json.dumps({"status": "OK"})
+        return res
+
+    def test_category_name_throws_errors_for_categories_that_cant_be_converted_to_filenames(self):
+        with self.app.test_request_context():
+            bad_category_names = make_fbs(
+                {
+                    "cat_A": pd.Series(["label_A"] * self.n_rows, dtype="category"),
+                    "cat/B": pd.Series(["label_B"] * self.n_rows, dtype="category"),
+                }
+            )
+            with self.assertRaises(AnnotationCategoryNameError):
+                self.annotation_put_fbs(bad_category_names)
+
+    def test_convert_to_pandas__converts_tiledb_to_pandas_df(self):
+        with self.app.test_request_context():
+            self.annotations.write_labels(self.df, self.data)
+            dataset_id = self.db.query([CellxGeneDataset], [CellxGeneDataset.name == self.data.get_location()])[0].id
+            annotation = self.db.query_for_most_recent(
+                Annotation,
+                [Annotation.user_id == self.user_id, Annotation.dataset_id == str(dataset_id)]
+            )
+            # retrieve tiledb array
+            df = tiledb.open(annotation.tiledb_uri)
+            self.assertEqual(type(df), tiledb.array.SparseArray)
+
+            # convert to pandas df
+            pandas_df = self.annotations.convert_to_pandas_df(df)
+            self.assertEqual(type(pandas_df), pd.DataFrame)
+
+    def test_write_labels_creates_a_dataset_if_it_doesnt_exist(self):
+        with self.app.test_request_context():
+
+            new_name = 'new_dataset/location'
+            self.data.get_location = MagicMock(return_value=new_name)
+            num_datasets = len(self.db.query([CellxGeneDataset]))
+            self.annotation_put_fbs(self.fbs)
+            more_datasets = len(self.db.query([CellxGeneDataset]))
+            self.assertGreater(more_datasets, num_datasets)
+
+            self.assertGreater(len(self.db.query([CellxGeneDataset], [CellxGeneDataset.name == new_name])), 0)
+
+    def test_write_labels_links_to_existing_dataset(self):
+        with self.app.test_request_context():
+            # add dataset to to db
+            self.annotation_put_fbs(self.fbs)
+
+            num_datasets = len(self.db.query([CellxGeneDataset]))
+
+            # create another annotation with the same dataset
+            self.annotation_put_fbs(self.fbs)
+
+            same_num_datasets = len(self.db.query([CellxGeneDataset]))
+
+            self.assertEqual(num_datasets, same_num_datasets)
+
+    def test_read_labels_returns_pandas_df(self):
+        with self.app.test_request_context():
+            self.annotation_put_fbs(self.fbs)
+            pandas_df = self.annotations.read_labels(self.data)
+            self.assertEqual(type(pandas_df), pd.DataFrame)
+
+    def test_read_labels_returns_df_matching_original(self):
+        with self.app.test_request_context():
+            self.annotation_put_fbs(self.fbs)
+            pandas_df = self.annotations.read_labels(self.data)
+
+            self.assertEqual(pandas_df.shape, (self.n_rows, 2))
+            self.assertEqual(set(pandas_df.columns), {"cat_A", "cat_B"})
+            self.assertTrue(self.data.original_obs_index.equals(pandas_df.index))
+            self.assertTrue(np.all(pandas_df["cat_A"] == ["label_A"] * self.n_rows))
+            self.assertTrue(np.all(pandas_df["cat_B"] == ["label_B"] * self.n_rows))
+
+    def test_error_checks(self):
+        # verify that the expected errors are generated
+        with self.app.test_request_context():
+            n_rows = self.data.get_shape()[0]
+            fbs_bad = make_fbs({"louvain": pd.Series(["undefined"] * n_rows, dtype="category")})
+
+            # ensure we catch attempt to overwrite non-writable data
+            with self.assertRaises(KeyError):
+                self.annotation_put_fbs(fbs_bad)
+
+    @patch('server.common.annotations.hosted_tiledb.current_app')
+    def test_write_labels_stores_df_as_tiledb_array(self, mock_user_id):
+        mock_user_id.auth.get_user_id.return_value = '1234'
+        self.annotations.write_labels(self.df, self.data)
+        # get uri
+        dataset_id = self.db.query([CellxGeneDataset], [CellxGeneDataset.name == self.data.get_location()])[0].id
+        annotation = self.db.query_for_most_recent(
+            Annotation,
+            [Annotation.user_id == '1234', Annotation.dataset_id == str(dataset_id)]
+        )
+
+        df = tiledb.open(annotation.tiledb_uri)
+        self.assertEqual(type(df), tiledb.array.SparseArray)
 
 
 class WritableAnnotationTest(unittest.TestCase):