Merge pull request #338 from vshekar/dewar-tree-faster-refresh

Added single call to get all dewar data and remove get() defaults
NSLS-II · Oct 16, 2023 · a5bd381 · a5bd381
2 parents 4116872 + 3d7c924
commit a5bd381
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 23 deletions.
diff --git a/db_lib.py b/db_lib.py
@@ -1,16 +1,15 @@
+import logging
 import os
-
 import time
-
-import six
-
 import uuid
+from collections import defaultdict
 
 import amostra.client.commands as acc
 import conftrak.client.commands as ccc
-from analysisstore.client.commands import AnalysisClient
 import conftrak.exceptions
-import logging
+import six
+from analysisstore.client.commands import AnalysisClient
+
 logger = logging.getLogger(__name__)
 
 #12/19 - Skinner inherited this from Hugo, who inherited it from Matt. Arman wrote the underlying DB and left BNL in 2018. 
@@ -486,6 +485,43 @@ def getContainerByID(container_id):
     c = getContainers(filters={'uid': container_id})[0]
     return c
 
+def get_dewar_tree_data(dewar_name, beamline):
+    """
+    returns all data required to show dewar tree data with minimum number of database accesses
+    """
+    dewar_data = getContainers(filters={"name": dewar_name, "owner": beamline})[0]
+
+    puck_ids = [
+        pid for pid in dewar_data.get("content", []) if pid
+    ]  # removes blank ids
+    pucks = getContainers(filters={"uid": {"$in": puck_ids}})
+
+    # Create a mega list of sample ids from puck information
+    sample_ids = [
+        sample_id
+        for puck in pucks
+        for sample_id in puck.get("content", [])
+        if sample_id
+    ]
+
+    # Get all sample info in one call
+    params = {"uid": {"$in": sample_ids}}
+    samples = sample_ref.find(as_document=False, **params)
+
+    # Get all request info in one call
+    params = {"sample": {"$in": sample_ids}, "state": "active"}
+    reqs = list(request_ref.find(**params))
+
+    # Assemble data into dictionaries
+    puck_data = {puck["uid"]: puck for puck in pucks}
+    sample_data = {sample["uid"]: sample for sample in samples}
+    request_data = defaultdict(list)
+    for req in reqs:
+        request_data[req["sample"]].append(req)
+
+    return dewar_data, puck_data, sample_data, request_data
+
+
 
 def getQueue(beamlineName):
     """

diff --git a/gui/dewar_tree.py b/gui/dewar_tree.py
@@ -1,16 +1,19 @@
+import getpass
 import logging
+import os
 import typing
-import os, getpass
+
+import requests
 from qtpy import QtCore, QtGui, QtWidgets
 from qtpy.QtCore import Qt
-import requests
+
 import daq_utils
 import db_lib
 from config_params import (
     DEWAR_SECTORS,
+    IS_STAFF,
     PUCKS_PER_DEWAR_SECTOR,
     SAMPLE_TIMER_DELAY,
-    IS_STAFF,
 )
 
 if typing.TYPE_CHECKING:
@@ -109,20 +112,17 @@ def set_mounted_sample(self, item):
     def refreshTreeDewarView(self):
         puck = ""
         self.model.clear()
-        dewarContents = db_lib.getContainerByName(
+        dewar_data, puck_data, sample_data, request_data = db_lib.get_dewar_tree_data(
             daq_utils.primaryDewarName, daq_utils.beamline
-        )["content"]
+        )
         parentItem = self.model.invisibleRootItem()
         for i, puck_id in enumerate(
-            dewarContents
+            dewar_data["content"]
         ):  # dewar contents is the list of puck IDs
             puck = ""
             puckName = ""
             if puck_id:
-                puck = containerDict.get(
-                    puck_id,
-                    containerDict.setdefault(puck_id, db_lib.getContainerByID(puck_id)),
-                )
+                puck = puck_data[puck_id]
                 puckName = puck["name"]
             sector, puck_pos = divmod(i, self.pucksPerDewarSector)
             index_s = f"{sector+1}{chr(puck_pos + ord('A'))}"
@@ -132,7 +132,9 @@ def refreshTreeDewarView(self):
             parentItem.appendRow(item)
             if puck != "" and puckName != "private":
                 puckContents = puck.get("content", [])
-                self.add_samples_to_puck_tree(puckContents, item, index_s)
+                self.add_samples_to_puck_tree(
+                    puckContents, item, index_s, sample_data, request_data
+                )
         self.setModel(self.model)
         self.model.itemChanged.connect(self.queueSelectedSample)
         if self.isExpanded:
@@ -142,7 +144,12 @@ def refreshTreeDewarView(self):
         self.scrollTo(self.currentIndex(), QtWidgets.QAbstractItemView.PositionAtCenter)
 
     def add_samples_to_puck_tree(
-        self, puckContents, parentItem: QtGui.QStandardItem, index_label
+        self,
+        puckContents,
+        parentItem: QtGui.QStandardItem,
+        index_label,
+        sample_data,
+        request_data,
     ):
         # Method will attempt to add samples to the puck. If you don't belong to the proposal,
         # it will not add samples and clear the puck information
@@ -159,10 +166,7 @@ def add_samples_to_puck_tree(
                 parentItem.appendRow(item)
                 continue
 
-            sample = sampleDict.get(
-                sample_id,
-                sampleDict.setdefault(sample_id, db_lib.getSampleByID(sample_id)),
-            )
+            sample = sample_data[sample_id]
 
             if not IS_STAFF and not self.is_proposal_member(sample["proposalID"]):
                 # If the user is not part of the proposal and is not staff, don't fill tree
@@ -196,7 +200,7 @@ def add_samples_to_puck_tree(
             if sample_id == self.parent.selectedSampleID:
                 logger.info("found " + str(self.parent.SelectedItemData))
                 selectedSampleIndex = self.model.indexFromItem(item)
-            sampleRequestList = db_lib.getRequestsBySampleID(sample_id)
+            sampleRequestList = request_data[sample_id]
             for request in sampleRequestList:
                 if not ("protocol" in request["request_obj"]):
                     continue