Merge pull request #11 from Nuwa-genomics/development

fixed tests

app/Dashboard.py

@@ -1,24 +1,18 @@
-import hashlib
 import streamlit as st
-import pandas as pd
 from models.WorkspaceModel import WorkspaceModel
 from models.AdataModel import AdataModel
 from state.AdataState import AdataState
 from utils.file_utils import *
-from datetime import datetime
 from database.database import *
 from database.schemas import schemas
 from sqlalchemy.orm import Session
-from pathlib import Path
 import shutil
 import zipfile
 import hashlib as hash
 from hashlib import sha256
-import time
 import os
 import scanpy as sc
 from database.db_export import *
-import pickle
 
 #create databases if not already present
 Base.metadata.create_all(engine)
@@ -49,6 +43,8 @@ class Dashboard:
     """
     def __init__(self):
         self.conn: Session = SessionLocal()
+        if "refresh_page" not in st.session_state:
+            st.session_state["refresh_page"] = False
         if not os.path.exists('/streamlit-volume/exported_workspaces'):
             os.mkdir('/streamlit-volume/exported_workspaces') 
         self.set_workspaces()
@@ -57,6 +53,7 @@ def __init__(self):
     def set_workspaces(self):
         workspaces = self.conn.query(schemas.Workspaces).all()
         st.session_state['workspaces'] = [WorkspaceModel(id=workspace.id, workspace_name=workspace.workspace_name, data_dir=workspace.data_dir, created=workspace.created, description=workspace.description) for workspace in workspaces]
+
 
 
     def _write_workspace_to_db(self):
@@ -73,12 +70,18 @@ def _write_workspace_to_db(self):
                 description = st.session_state.ti_new_workspace_desc,
                 data_dir = path
             )
-            self.conn.add(new_workspace)
-            self.conn.commit()
-            self.conn.refresh(new_workspace)
-            st.toast("Successfully created workspace", icon='✅')
+            exists = self.conn.query(schemas.Workspaces).filter_by(workspace_name=st.session_state.ti_new_workspace_name).first() is not None
+            if not exists:
+                self.conn.add(new_workspace)
+                self.conn.commit()
+                self.conn.refresh(new_workspace)
+                st.toast("Successfully created workspace", icon='✅')
+                return 0
+            else:
+                return -1
+
         except Exception as e:
-            st.error(e)
+            print(e)
             st.toast("Failed to create workspace", icon="❌")
 
 
@@ -114,15 +117,12 @@ def new_workspace(self):
         .. image:: https://raw.githubusercontent.com/nuwa-genomics/Nuwa/main/docs/assets/images/screenshots/new_workspace.png
         """
         with st.sidebar:
-            with st.form(key="new_workspace_form"):
+            with st.container(border=True):
                 st.subheader("Create New Workspace")
                 st.text_input(label="Name", key="ti_new_workspace_name")
                 st.text_input(label="Description", key="ti_new_workspace_desc")
                 col1, _, _ = st.columns(3)
-                save_btn = col1.form_submit_button(label="Save", use_container_width=True)
-                if save_btn:
-                    self._write_workspace_to_db()
-                    self.set_workspaces()
+                col1.button(label="Save", use_container_width=True, on_click=self._write_workspace_to_db, key="btn:dashboard:new_workspace")
 
 
     def import_workspace(self):
@@ -255,18 +255,17 @@ def delete_workspace(self):
             st.toast("Successfully deleted workspace", icon="✅")
         except Exception as e:
             st.toast("Failed to delete workspace", icon="❌")
+            print(e)
 
 
-
+    
     def draw_page(self):
         st.title("Workspaces")
-
-        subheader_container = st.container()
+
         col1, col2, col3, col4 = st.columns(4, gap="large")
         columns = [col1, col2, col3, col4]
 
-        if len(st.session_state.workspaces) > 0:
-            subheader_container.subheader("Select a workspace")
+        if len(st.session_state["workspaces"]) > 0:
             for i, workspace in enumerate(st.session_state.workspaces): 
                 with columns[i % 4]:
                     button = st.button(label=workspace.workspace_name, key=f"btn_workspace_{workspace.id}")
@@ -311,7 +310,4 @@ def draw_page(self):
                         <div style='font-size: 16px; color: rgba(255, 255, 255, 0.4)'>Nuwa v{os.getenv('NUWA_VERSION')}</div>
                         </div>""", unsafe_allow_html=True)
 
-dashboard = Dashboard()
-
-
-
+dashboard = Dashboard()

app/components/sidebar.py

@@ -12,7 +12,8 @@
 from utils.Gene_info import Gene_info
 from utils.species import *
 import numpy as np
-
+from state.StateManager import StateManager
+from models.ErrorMessage import ErrorMessage, WarningMessage
 
 
 class Sidebar:
@@ -21,6 +22,7 @@ class Sidebar:
     """
     def __init__(self):
         self.conn: Session = SessionLocal()
+        self.state_manager = StateManager()
 
     def show_preview(self, integrate=False):
         """
@@ -44,13 +46,14 @@ def show_preview(self, integrate=False):
                     st.markdown(f"""<h3 style='font-size: 16px; color: rgba(255, 255, 255, 0.75)'>{st.session_state.adata_target.adata_name}</h3>""", unsafe_allow_html=True)
                     st.markdown(f"""<p style='font-size: 14px; color: rgba(255, 255, 255, 0.75)'>{st.session_state.adata_target.adata}</p>""", unsafe_allow_html=True)
                 else:
-                    st.markdown(f"""<p style='font-size: 14px; color: rgba(255, 255, 255, 0.75)'>{st.session_state.adata_state.current.adata}</p>""", unsafe_allow_html=True)
+                    st.markdown(f"""<p style='font-size: 14px; color: rgba(255, 255, 255, 0.75)'>{self.state_manager.get_current_adata()}</p>""", unsafe_allow_html=True)
 
     def delete_experiment_btn(self):
         with st.sidebar:
             delete_btn = st.button(label="🗑️ Delete Experiment", use_container_width=True, key="btn_delete_adata")
             if delete_btn:
-                st.session_state.adata_state.delete_record(adata_name=st.session_state.sb_adata_selection)
+                self.state_manager.adata_state() \
+                .delete_record(adata_name=st.session_state.sb_adata_selection)
 
 
     def export_script(self):
@@ -93,15 +96,18 @@ def notes(self):
         .. image:: https://raw.githubusercontent.com/nuwa-genomics/Nuwa/main/docs/assets/images/screenshots/notes.png
         """
         with st.sidebar:
-
-            notes = st.session_state.adata_state.load_adata(workspace_id=st.session_state.current_workspace.id, adata_name=st.session_state.sb_adata_selection).notes
-            display_notes = notes if notes != None else ""
-            notes_ta = st.text_area(label="Notes", placeholder="Notes", value=display_notes, key="sidebar_notes")
             try:
-                st.session_state.adata_state.current.notes = notes_ta
-                st.session_state.adata_state.update_record()
+                load_adata = self.state_manager.adata_state() \
+                .load_adata(workspace_id=st.session_state.current_workspace.id, adata_name=st.session_state.sb_adata_selection)
+
+                notes = load_adata.notes
+                display_notes = notes if notes != None else ""
+                notes_ta = st.text_area(label="Notes", placeholder="Notes", value=display_notes, key="sidebar_notes")
+
+                self.state_manager.adata_state().current.notes = notes_ta
+                self.state_manager.adata_state().update_record()
             except Exception as e:
-                st.toast("Notes failed to save", icon="❌")
+                st.toast(ErrorMessage.NOTES_FAILED_TO_SAVE.value, icon="❌")
                 print("Error: ", e)
 
     def get_adata(self, adataList, name) -> AdataModel:
@@ -135,7 +141,11 @@ def add_experiment(self):
                 try:
                     st.subheader("Create New Adata")
                     st.text_input(label="Name", key="ti_new_adata_name")
-                    st.selectbox(label="Dataset", key="sb_new_experiment_dataset", options=st.session_state.adata_state.get_adata_options())
+                    adata_options = self.state_manager.adata_state().get_adata_options()
+                    if isinstance(adata_options, ErrorMessage):
+                        st.toast(adata_options.value, icon="❌")
+                        return
+                    st.selectbox(label="Dataset", key="sb_new_experiment_dataset", options=adata_options)
                     st.text_input(label="Notes", key="ti_new_adata_notes")
                     st.button(label="Save", on_click=self.write_adata, key="btn_add_adata")
                 except Exception as e:
@@ -166,10 +176,11 @@ def download_adata(self):
                     st.text_input(label="Download directory", value=os.path.join(os.getenv('WORKDIR'), 'downloads', st.session_state.sb_adata_selection), key="ti_save_adata_dir")
                     save_adata_btn = st.button(label="Download", key="btn_save_adata")
                     if save_adata_btn:
-                        selected_adata = st.session_state.adata_state.load_adata(workspace_id=st.session_state.current_workspace.id, adata_name=st.session_state.sb_adata_selection)
+                        selected_adata = self.state_manager.adata_state() \
+                        .load_adata(workspace_id=st.session_state.current_workspace.id, adata_name=st.session_state.sb_adata_selection)
 
-                        if not selected_adata:
-                            st.toast("Couldn't find selected adata to save", icon="❌")
+                        if isinstance(selected_adata, ErrorMessage):
+                            st.toast(selected_adata.value, icon="❌")
                         else:
                             download_path = os.path.join(os.getenv('WORKDIR'), 'downloads', st.session_state.sb_adata_selection)
                             if not os.path.exists(download_path):
@@ -209,7 +220,7 @@ def write_adata(self):
             name = st.session_state.ti_new_adata_name
             notes = st.session_state.ti_new_adata_notes
 
-            st.session_state.adata_state.insert_record(AdataModel(
+            self.state_manager.adata_state().insert_record(AdataModel(
                 work_id=st.session_state.current_workspace.id,
                 adata_name=name,
                 filename=os.path.join(os.getenv('WORKDIR'), "adata", f"{name}.h5ad"),
@@ -247,14 +258,14 @@ def gene_format():
         """
         try:
             format = ""
-            for var in st.session_state.adata_state.current.adata.var_names[:5]:
+            for var in StateManager().get_current_adata().var_names[:5]:
                 if not var.startswith('ENS'):
                     format = "gene_symbol"
                 else:
                     format = "ensembl"
 
             def change_gene_format():
-                gene_info = Gene_info(st.session_state.adata_state.current.adata)
+                gene_info = Gene_info(StateManager().get_current_adata())
                 if gene_info.fail == -1:
                     st.session_state["toggle_gene_format"] = not st.session_state["toggle_gene_format"]
                     return -1
@@ -263,7 +274,7 @@ def change_gene_format():
                 else:
                     gene_info.convert_enseml_to_symbols() #change format to gene symbols
 
-                st.session_state.adata_state.current.adata.var = gene_info.adata.var
+                StateManager().get_current_adata().var = gene_info.adata.var
 
                 if st.session_state.toggle_gene_format:
                     st.toast("Changed format to Ensembl IDs", icon='✅')
@@ -290,18 +301,40 @@ def show_version(self):
     def show(self, integrate = False):
         with st.sidebar:
             def set_adata():
-                if st.session_state.adata_state.switch_adata(st.session_state.sb_adata_selection) == -1:
-                    st.error("Couldn't switch adata")
+                switch_adata_result = self.state_manager.adata_state().switch_adata(st.session_state.sb_adata_selection)
+                if isinstance(switch_adata_result, ErrorMessage):
+                    st.toast(switch_adata_result.value, icon="❌")
 
 
             if integrate:
-                st.selectbox(label="Current Experiment (reference adata):", options=st.session_state.adata_state.get_adata_options(), key="sb_adata_selection", on_change=set_adata, index=st.session_state.adata_state.get_index_of_current())
-                st.selectbox(label="Integrate into:", options=st.session_state.adata_state.get_adata_options(), key="sb_adata_selection_target")
+                options = self.state_manager.adata_state().get_adata_options()
+                index = self.state_manager.adata_state().get_index_of_current()
+                if isinstance(options, ErrorMessage):
+                    st.toast(options.value, icon="❌")
+                    return
+                if isinstance(index, ErrorMessage):
+                    st.toast(options.value, icon="❌")
+                    return
+                st.selectbox(
+                    label="Current Experiment (reference adata):", 
+                    options=options, 
+                    key="sb_adata_selection", 
+                    on_change=set_adata, 
+                    index=index
+                )
+                st.selectbox(
+                    label="Integrate into:", 
+                    options=options, 
+                    key="sb_adata_selection_target"
+                )
 
-
                 #set integrate adata
-                st.session_state['adata_ref']: AdataModel = st.session_state.adata_state.current.copy()
-                st.session_state['adata_target']: AdataModel = st.session_state.adata_state.load_adata(st.session_state.current_workspace.id, st.session_state.sb_adata_selection_target).copy()
+
+                st.session_state['adata_ref'] = self.state_manager.adata_state().current
+                st.session_state['adata_target'] = self.state_manager.adata_state().load_adata(
+                    workspace_id=st.session_state.current_workspace.id, 
+                    adata_name=st.session_state.sb_adata_selection_target
+                )
 
                 #sync genes
                 #test var names
@@ -339,14 +372,21 @@ def set_adata():
                     empty.markdown(f"""<p style='font-size: 15px; color: rgba(255, 255, 255, 0.75)'>Current experiment {len(st.session_state.adata_ref.adata.var_names)} genes</p><p style='font-size: 15px; color: rgba(255, 255, 255, 0.75)'>Target experiment {len(st.session_state.adata_target.adata.var_names)} genes</p>""", unsafe_allow_html=True)
                 else:
                     #set integrate adata
-                    st.session_state['adata_ref']: AdataModel = st.session_state.adata_state.current
-                    st.session_state['adata_target']: AdataModel = st.session_state.adata_state.load_adata(st.session_state.current_workspace.id, st.session_state.sb_adata_selection_target)
+                    st.session_state['adata_ref'] = self.state_manager.adata_state().current
+                    st.session_state['adata_target'] = self.state_manager.adata_state().load_adata(
+                        workspace_id=st.session_state.current_workspace.id, 
+                        adata_name=st.session_state.sb_adata_selection_target
+                    )
                     empty.markdown(f"""<p style='font-size: 15px; color: rgba(255, 255, 255, 0.75)'>Current experiment {len(st.session_state.adata_ref.adata.var_names)} genes</p><p style='font-size: 15px; color: rgba(255, 255, 255, 0.75)'>Target experiment {len(st.session_state.adata_target.adata.var_names)} genes</p>""", unsafe_allow_html=True)
 
                 st.divider()
 
             else:
-                st.selectbox(label="Current Experiment:", options=st.session_state.adata_state.get_adata_options(), key="sb_adata_selection", on_change=set_adata, index=st.session_state.adata_state.get_index_of_current())
+                options = self.state_manager.adata_state().get_adata_options()
+                if isinstance(options, ErrorMessage):
+                    st.toast(options.value, icon="❌")
+                    return
+                st.selectbox(label="Current Experiment:", options=options, key="sb_adata_selection", on_change=set_adata, index=self.state_manager.adata_state().get_index_of_current())
 
                 Sidebar.species()
 

app/database/database.py

@@ -20,7 +20,7 @@ def get_db():
     finally:
         db.close()
 
-engine = create_engine(SQLALCHEMY_DATABASE_URL)
+engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_size=60, max_overflow=20)
 
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 

app/models/ErrorMessage.py

@@ -0,0 +1,15 @@
+from enum import Enum
+import streamlit as st
+
+class ErrorMessage(Enum):
+    ADATA_NOT_FOUND = "Couldn't find adata"
+    AT_LEAST_ONE_EXPERIMENT_REQUIRED = "Workspace must contain at least one experiment"
+    CANNOT_DELETE_EXPERIMENT = "Couldn't delete experiment"
+    CANNOT_UPDATE_EXPERIMENT = "Couldn't update experiment"
+    CANNOT_ADD_EXPERIMENT = "Couldn't add experiment"
+    FILE_DOES_NOT_EXIST = "File does not exist"
+    NOTES_FAILED_TO_SAVE = "Notes failed to save"
+
+class WarningMessage(Enum):
+    DATASET_ALREADY_EXISTS = "Dataset already exists in workspace, using original."
+

app/models/ScriptModel.py

@@ -13,4 +13,4 @@ class ScriptModel(BaseModel):
 class Language(Enum):
     python = "python"
     R = "R"
-    ALL_SUPPORTED = [python, R]
+    ALL_SUPPORTED = ["python", "R"]