Dunitz/db setup (#1619)

* initial database setup

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)