Dunitz/db setup (#1619)

* initial database setup

Makefile

@@ -60,6 +60,11 @@ smoke-test:
 smoke-test-annotations:
 	cd client && $(MAKE) smoke-test-annotations
 
+.PHONY: test-db
+test-db:
+	cd server && $(MAKE) test-db
+
+
 # FORMATTING CODE
 
 .PHOHY: fmt

server/Makefile

@@ -7,11 +7,35 @@ clean:
 	rm -f common/web/csp-hashes.json
 
 .PHONY: unit-test
-unit-test:
+unit-test: create-test-db
 	PYTHONWARNINGS=ignore:ResourceWarning coverage run \
 		--source=app,cli,common,compute,converters,data_anndata,data_common,data_cxg \
 		--omit=.coverage,data_common/fbs/NetEncoding,venv \
 		-m unittest discover \
 		--start-directory test/ \
 		--top-level-directory ../ \
-		--verbose
+		--verbose; test_result=$$?; \
+	$(MAKE) clean-test-db; \
+	exit $$test_result \
+
+
+.PHONY: test-db
+test-db: create-test-db
+	PYTHONWARNINGS=ignore:ResourceWarning coverage run \
+		--source=app,cli,common,compute,converters,data_anndata,data_common,data_cxg \
+		--omit=.coverage,data_common/fbs/NetEncoding,venv \
+		-m unittest discover \
+		--start-directory test/test_database \
+		--top-level-directory ../ \
+		--verbose; test_result=$$?; \
+	$(MAKE) clean-test-db; \
+	exit $$test_result
+
+.PHONY: create-test-db
+create-test-db:
+	-docker run -d -p 5432:5432 --name test_db -e POSTGRES_PASSWORD=test_pw postgres
+
+.PHONY: clean-test-db
+clean-test-db:
+	-docker stop test_db
+	-docker rm test_db

server/db/cellxgene_orm.py

@@ -0,0 +1,60 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    DateTime,
+    ForeignKey,
+    String,
+)
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship
+
+Base = declarative_base()
+
+
+class CellxGeneUser(Base):
+    """
+    A registered CellxGene user.
+    Links a user to their annotations
+    """
+
+    __tablename__ = "cxguser"
+
+    id = Column(String, primary_key=True)
+    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
+    updated_at = Column(DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+    # Relationships
+    annotations = relationship("Annotation", back_populates="cxguser")
+
+
+class Annotation(Base):
+    """
+    An annotation is a link between a user, a dataset and tiledb dataframe. A user can have multiple annotations for a
+    dataset, the most recent annotation (based on created_at) will be the default returned when queried
+    """
+
+    __tablename__ = "annotation"
+
+    id = Column(String, primary_key=True)
+    tiledb_uri = Column(String)
+    user_id = Column(String, ForeignKey("cxguser.id"), nullable=False)
+    dataset_id = Column(String, ForeignKey("cxgdataset.id"), nullable=False)
+    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
+
+    # Relationships
+    cxguser = relationship("CellxGeneUser", back_populates="annotations")
+    dataset = relationship("CellxGeneDataset", back_populates="annotations")
+
+
+class CellxGeneDataset(Base):
+    """
+    Datasets refer to cellxgene datasets stored in tiledb
+    """
+
+    __tablename__ = "cxgdataset"
+
+    id = Column(String, primary_key=True)
+    name = Column(String)
+    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
+    annotations = relationship("Annotation", back_populates="dataset")

server/db/create_db.py

@@ -0,0 +1,18 @@
+"""
+Drops and recreates all tables for local testing according to cellxgene_orm.py
+"""
+from sqlalchemy import create_engine
+
+from server.db.cellxgene_orm import Base
+
+
+def create_db(database_uri: str = "postgresql://postgres:test_pw@localhost:5432"):
+    engine = create_engine(database_uri)
+    print("Dropping tables")
+    Base.metadata.drop_all(engine)
+    print("Recreating tables")
+    Base.metadata.create_all(engine)
+
+
+if __name__ == "__main__":
+    create_db()

server/db/db_utils.py

@@ -0,0 +1,43 @@
+import typing
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from server.db.cellxgene_orm import Base
+
+
+class DbUtils:
+    def __init__(self, database_uri: str = "postgresql://postgres:test_pw@localhost:5432"):
+        self.session = DBSessionMaker(database_uri).session()
+        self.engine = self.session.get_bind()
+
+    def get(self, table: Base, entity_id: typing.Union[str, typing.Tuple[str]]) -> typing.Union[Base, None]:
+        """
+        Query a table row by its primary key
+        :param table: SQLAlchemy Table to query
+        :param entity_id: Primary key of desired row
+        :return: SQLAlchemy Table object, None if not found
+        """
+        return self.session.query(table).get(entity_id)
+
+    def query(self, table_args: typing.List[Base], filter_args: typing.List[bool] = None) -> typing.List[Base]:
+        """
+        Query the database using the current DB session
+        :param table_args: List of SQLAlchemy Tables to query/join
+        :param filter_args: List of SQLAlchemy filter conditions
+        :return: List of SQLAlchemy query response objects
+        """
+        return (
+            self.session.query(*table_args).filter(*filter_args).all()
+            if filter_args
+            else self.session.query(*table_args).all()
+        )
+
+
+class DBSessionMaker:
+    def __init__(self, database_uri):
+        self.engine = create_engine(database_uri, connect_args={"connect_timeout": 5})
+        self.session_maker = sessionmaker(bind=self.engine)
+
+    def session(self, **kwargs):
+        return self.session_maker(**kwargs)

server/eb/app.py

@@ -31,7 +31,7 @@ except Exception:
     sys.exit(1)
 
 
-def get_flask_secret_key(region_name, secret_name):
+def get_secret_key(region_name, secret_name, secret_key):
     session = boto3.session.Session()
     client = session.client(service_name="secretsmanager", region_name=region_name)
 
@@ -40,7 +40,7 @@ def get_flask_secret_key(region_name, secret_name):
         if "SecretString" in get_secret_value_response:
             var = get_secret_value_response["SecretString"]
             secret = json.loads(var)
-            return secret.get("flask_secret_key")
+            return secret.get(secret_key)
     except Exception:
         logging.critical("Caught exception during get_secret_key", exc_info=True)
         sys.exit(1)
@@ -173,7 +173,7 @@ try:
             logging.error("Could not determine the AWS Secret Manager region")
             sys.exit(1)
 
-        flask_secret_key = get_flask_secret_key(secret_region_name, secret_name)
+        flask_secret_key = get_secret_key(secret_region_name, secret_name, 'flask_secret_key')
         app_config.update_server_config(app__flask_secret_key=flask_secret_key)
 
     # features are unsupported in the current hosted server

server/requirements.txt

@@ -15,9 +15,11 @@ numba>=0.49.1
 numpy>=1.16.0
 packaging>=20.0
 pandas>=0.24.2
+psycopg2==2.7.7
 PyYAML>=5.3
 scipy>=1.3.0
 requests>=2.22.0
+sqlalchemy>=1.3.18
 tiledb>=0.5.9,>=0.6.2
 s3fs>=0.4.2
 gunicorn>=20.0.4

server/test/fixtures/database/__init__.py

@@ -0,0 +1,87 @@
+import string
+import random
+
+
+from sqlalchemy import func
+
+from server.db.cellxgene_orm import CellxGeneUser, CellxGeneDataset, Annotation, Base
+from server.db.create_db import create_db
+from server.db.db_utils import DbUtils
+
+
+class TestDatabase:
+    def __init__(self):
+        local_db_uri = "postgresql://postgres:test_pw@localhost:5432"
+        create_db(local_db_uri)
+        self.db = DbUtils(local_db_uri)
+        self._populate_test_data()
+        self._populate_test_data_many()
+
+    def _populate_test_data(self):
+        self._create_test_user()
+        self._create_test_dataset()
+        self._create_test_annotation()
+
+    def _populate_test_data_many(self):
+        self._create_test_users()
+        self._create_test_datasets()
+        self._create_test_annotations()
+
+    def _create_test_user(self):
+        user = CellxGeneUser(id="test_user_id")
+        self.db.session.add(user)
+        self.db.session.commit()
+
+    def _create_test_dataset(self):
+        dataset = CellxGeneDataset(
+            id="test_dataset_id",
+            name="test_dataset",
+        )
+        self.db.session.add(dataset)
+        self.db.session.commit()
+
+    def _create_test_annotation(self):
+        annotation = Annotation(
+            id="test_annotation_id",
+            tiledb_uri="tiledb_uri",
+            user_id="test_user_id",
+            dataset_id="test_dataset_id"
+        )
+        self.db.session.add(annotation)
+        self.db.session.commit()
+
+    @staticmethod
+    def get_random_string():
+        letters = string.ascii_lowercase
+        return ''.join(random.choice(letters) for i in range(12))
+
+    def _create_test_users(self, user_count: int = 10):
+        users = []
+        for i in range(user_count):
+            users.append(CellxGeneUser(id=self.get_random_string()))
+        self.db.session.add_all(users)
+        self.db.session.commit()
+
+    def _create_test_datasets(self, dataset_count: int = 10):
+        datasets = []
+        for i in range(dataset_count):
+            datasets.append(CellxGeneDataset(id=self.get_random_string(), name=self.get_random_string()))
+        self.db.session.add_all(datasets)
+        self.db.session.commit()
+
+    def order_by_random(self, table: Base):
+        return self.db.session.query(table).order_by(func.random()).first()
+
+    def _create_test_annotations(self, annotation_count: int = 10):
+        annotations = []
+        for i in range(annotation_count):
+            dataset = self.order_by_random(CellxGeneDataset)
+            user = self.order_by_random(CellxGeneUser)
+            annotations.append(Annotation(
+                id=self.get_random_string(),
+                tiledb_uri=self.get_random_string(),
+                user_id=user.id,
+                dataset_id=dataset.id
+            ))
+        self.db.session.add_all(annotations)
+        self.db.session.commit()

server/test/test_database/test_database.py

@@ -0,0 +1,34 @@
+import unittest
+from server.db.cellxgene_orm import CellxGeneUser, CellxGeneDataset, Annotation
+from server.db.db_utils import DbUtils
+from server.test.fixtures.database import TestDatabase
+
+
+class AppConfigTest(unittest.TestCase):
+    db = DbUtils("postgresql://postgres:test_pw@localhost:5432")
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        TestDatabase()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        del cls.db
+
+    def test_user_creation(self):
+        one_user = self.db.get(table=CellxGeneUser, entity_id='test_user_id')
+        self.assertEqual(one_user.id, 'test_user_id')
+        user_count = self.db.session.query(CellxGeneUser).count()
+        self.assertGreater(user_count, 10)
+
+    def test_dataset_creation(self):
+        one_dataset = self.db.get(table=CellxGeneDataset, entity_id='test_dataset_id')
+        self.assertEqual(one_dataset.id, 'test_dataset_id')
+        dataset_count = self.db.session.query(CellxGeneDataset).count()
+        self.assertGreater(dataset_count, 10)
+
+    def test_annotation_creation(self):
+        one_annotation = self.db.get(table=Annotation, entity_id='test_annotation_id')
+        self.assertEqual(one_annotation.id, 'test_annotation_id')
+        annotation_count = self.db.session.query(Annotation).count()
+        self.assertGreater(annotation_count, 10)